From 26adb4dbae6be033c47014217a3283c394880b26 Mon Sep 17 00:00:00 2001 From: Nick Fitzgerald Date: Tue, 1 Aug 2023 13:12:25 -0700 Subject: [PATCH 01/11] Wasmtime: Rename `IndexAllocator` to `ModuleAffinityIndexAllocator` We will have multiple kinds of index allocators soon, so clarify which one this is. --- .../runtime/src/instance/allocator/pooling.rs | 10 +++++----- .../allocator/pooling/index_allocator.rs | 18 +++++++++--------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/crates/runtime/src/instance/allocator/pooling.rs b/crates/runtime/src/instance/allocator/pooling.rs index c35cf6ffd2e4..2ee511c6e76b 100644 --- a/crates/runtime/src/instance/allocator/pooling.rs +++ b/crates/runtime/src/instance/allocator/pooling.rs @@ -21,7 +21,7 @@ use wasmtime_environ::{ }; mod index_allocator; -use index_allocator::{IndexAllocator, SlotId}; +use index_allocator::{ModuleAffinityIndexAllocator, SlotId}; cfg_if::cfg_if! { if #[cfg(windows)] { @@ -375,7 +375,7 @@ struct StackPool { stack_size: usize, max_instances: usize, page_size: usize, - index_allocator: IndexAllocator, + index_allocator: ModuleAffinityIndexAllocator, async_stack_zeroing: bool, async_stack_keep_resident: usize, } @@ -427,7 +427,7 @@ impl StackPool { // Note that `max_unused_warm_slots` is set to zero since stacks // have no affinity so there's no need to keep intentionally unused // warm slots around. - index_allocator: IndexAllocator::new(config.limits.count, 0), + index_allocator: ModuleAffinityIndexAllocator::new(config.limits.count, 0), }) } @@ -572,7 +572,7 @@ impl Default for PoolingInstanceAllocatorConfig { pub struct PoolingInstanceAllocator { instance_size: usize, max_instances: usize, - index_allocator: IndexAllocator, + index_allocator: ModuleAffinityIndexAllocator, memories: MemoryPool, tables: TablePool, linear_memory_keep_resident: usize, @@ -596,7 +596,7 @@ impl PoolingInstanceAllocator { Ok(Self { instance_size: round_up_to_pow2(config.limits.size, mem::align_of::()), max_instances, - index_allocator: IndexAllocator::new(config.limits.count, config.max_unused_warm_slots), + index_allocator: ModuleAffinityIndexAllocator::new(config.limits.count, config.max_unused_warm_slots), memories: MemoryPool::new(&config.limits, tunables)?, tables: TablePool::new(&config.limits)?, linear_memory_keep_resident: config.linear_memory_keep_resident, diff --git a/crates/runtime/src/instance/allocator/pooling/index_allocator.rs b/crates/runtime/src/instance/allocator/pooling/index_allocator.rs index 94a001f4e6b1..a82dcc2986d4 100644 --- a/crates/runtime/src/instance/allocator/pooling/index_allocator.rs +++ b/crates/runtime/src/instance/allocator/pooling/index_allocator.rs @@ -17,7 +17,7 @@ impl SlotId { } #[derive(Debug)] -pub struct IndexAllocator(Mutex); +pub struct ModuleAffinityIndexAllocator(Mutex); #[derive(Debug)] struct Inner { @@ -115,10 +115,10 @@ enum AllocMode { AnySlot, } -impl IndexAllocator { +impl ModuleAffinityIndexAllocator { /// Create the default state for this strategy. pub fn new(max_instances: u32, max_unused_warm_slots: u32) -> Self { - IndexAllocator(Mutex::new(Inner { + ModuleAffinityIndexAllocator(Mutex::new(Inner { last_cold: 0, max_unused_warm_slots, unused_warm_slots: 0, @@ -410,13 +410,13 @@ impl List { #[cfg(test)] mod test { - use super::{IndexAllocator, SlotId}; + use super::{ModuleAffinityIndexAllocator, SlotId}; use crate::CompiledModuleIdAllocator; #[test] fn test_next_available_allocation_strategy() { for size in 0..20 { - let state = IndexAllocator::new(size, 0); + let state = ModuleAffinityIndexAllocator::new(size, 0); for i in 0..size { assert_eq!(state.alloc(None).unwrap().index(), i as usize); } @@ -429,7 +429,7 @@ mod test { let id_alloc = CompiledModuleIdAllocator::new(); let id1 = id_alloc.alloc(); let id2 = id_alloc.alloc(); - let state = IndexAllocator::new(100, 100); + let state = ModuleAffinityIndexAllocator::new(100, 100); let index1 = state.alloc(Some(id1)).unwrap(); assert_eq!(index1.index(), 0); @@ -483,7 +483,7 @@ mod test { let id = id_alloc.alloc(); for max_unused_warm_slots in [0, 1, 2] { - let state = IndexAllocator::new(100, max_unused_warm_slots); + let state = ModuleAffinityIndexAllocator::new(100, max_unused_warm_slots); let index1 = state.alloc(Some(id)).unwrap(); let index2 = state.alloc(Some(id)).unwrap(); @@ -504,7 +504,7 @@ mod test { let ids = std::iter::repeat_with(|| id_alloc.alloc()) .take(10) .collect::>(); - let state = IndexAllocator::new(1000, 1000); + let state = ModuleAffinityIndexAllocator::new(1000, 1000); let mut allocated: Vec = vec![]; let mut last_id = vec![None; 1000]; @@ -548,7 +548,7 @@ mod test { let id1 = id_alloc.alloc(); let id2 = id_alloc.alloc(); let id3 = id_alloc.alloc(); - let state = IndexAllocator::new(10, 2); + let state = ModuleAffinityIndexAllocator::new(10, 2); // Set some slot affinities assert_eq!(state.alloc(Some(id1)), Some(SlotId(0))); From aa0670128e4ccdb50da13c130273766bc6b4e1a6 Mon Sep 17 00:00:00 2001 From: Nick Fitzgerald Date: Wed, 9 Aug 2023 11:02:49 -0700 Subject: [PATCH 02/11] Wasmtime: Introduce a simple index allocator This will be used in future commits refactoring the pooling allocator. --- .../allocator/pooling/index_allocator.rs | 32 ++++++++++++++++--- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/crates/runtime/src/instance/allocator/pooling/index_allocator.rs b/crates/runtime/src/instance/allocator/pooling/index_allocator.rs index a82dcc2986d4..43433185a3dd 100644 --- a/crates/runtime/src/instance/allocator/pooling/index_allocator.rs +++ b/crates/runtime/src/instance/allocator/pooling/index_allocator.rs @@ -5,10 +5,10 @@ use std::collections::hash_map::{Entry, HashMap}; use std::mem; use std::sync::Mutex; -/// A slot index. The job of this allocator is to hand out these -/// indices. +/// A slot index. #[derive(Hash, Clone, Copy, Debug, PartialEq, Eq)] pub struct SlotId(pub u32); + impl SlotId { /// The index of this slot. pub fn index(self) -> usize { @@ -16,6 +16,30 @@ impl SlotId { } } +/// A simple index allocator. +/// +/// This index allocator doesn't do any module affinity or anything like that, +/// however it is built on top of the `ModuleAffinityIndexAllocator` to save +/// code (and code size). +#[derive(Debug)] +pub struct SimpleIndexAllocator(ModuleAffinityIndexAllocator); + +impl SimpleIndexAllocator { + pub fn new(capacity: u32) -> Self { + SimpleIndexAllocator(ModuleAffinityIndexAllocator::new(capacity, 0)) + } + + pub fn alloc(&self) -> Option { + self.0.alloc(None) + } + + pub(crate) fn free(&self, index: SlotId) { + self.0.free(index); + } +} + +/// An index allocator that has configurable affinity between slots and modules +/// so that slots are often reused for the same module again. #[derive(Debug)] pub struct ModuleAffinityIndexAllocator(Mutex); @@ -117,13 +141,13 @@ enum AllocMode { impl ModuleAffinityIndexAllocator { /// Create the default state for this strategy. - pub fn new(max_instances: u32, max_unused_warm_slots: u32) -> Self { + pub fn new(capacity: u32, max_unused_warm_slots: u32) -> Self { ModuleAffinityIndexAllocator(Mutex::new(Inner { last_cold: 0, max_unused_warm_slots, unused_warm_slots: 0, module_affine: HashMap::new(), - slot_state: (0..max_instances).map(|_| SlotState::UnusedCold).collect(), + slot_state: (0..capacity).map(|_| SlotState::UnusedCold).collect(), warm: List::default(), })) } From 5bc63dd0f3f0d212822f8a0ec52e8cb33433b1ea Mon Sep 17 00:00:00 2001 From: Nick Fitzgerald Date: Wed, 9 Aug 2023 11:13:02 -0700 Subject: [PATCH 03/11] Wasmtime: refactor the pooling allocator for components We used to have one index allocator, an index per instance, and give out N tables and M memories to every instance regardless how many tables and memories they need. Now we have an index allocator for memories and another for tables. An instance isn't associated with a single instance, each of its memories and tables have an index. We allocate exactly as many tables and memories as the instance actually needs. Ultimately, this gives us better component support, where a component instance might have varying numbers of internal tables and memories. Additionally, you can now limit the number of tables, memories, and core instances a single component can allocate from the pooling allocator, even if there is the capacity for that many available. This is to give embedders tools to limit individual component instances and prevent them from hogging too much of the pooling allocator's resources. --- RELEASES.md | 45 + benches/instantiation.rs | 2 +- crates/fuzzing/src/generators/config.rs | 40 +- .../fuzzing/src/generators/pooling_config.rs | 104 +- crates/jit/src/instantiate.rs | 2 +- crates/runtime/src/instance.rs | 54 +- crates/runtime/src/instance/allocator.rs | 414 ++++-- .../src/instance/allocator/on_demand.rs | 133 +- .../runtime/src/instance/allocator/pooling.rs | 1272 ++++------------- .../allocator/pooling/index_allocator.rs | 111 +- .../instance/allocator/pooling/memory_pool.rs | 439 ++++++ .../src/instance/allocator/pooling/simple.rs | 24 + .../instance/allocator/pooling/stack_pool.rs | 237 +++ .../instance/allocator/pooling/table_pool.rs | 210 +++ crates/runtime/src/lib.rs | 4 +- crates/wasi/src/preview2/preview1.rs | 12 +- crates/wasmtime/src/component/component.rs | 12 +- crates/wasmtime/src/component/instance.rs | 20 +- crates/wasmtime/src/config.rs | 285 +++- crates/wasmtime/src/instance.rs | 2 +- crates/wasmtime/src/module.rs | 6 +- crates/wasmtime/src/store.rs | 46 +- crates/wasmtime/src/trampoline.rs | 7 +- crates/wasmtime/src/trampoline/memory.rs | 118 +- fuzz/fuzz_targets/instantiate-many.rs | 2 +- tests/all/async_functions.rs | 13 +- tests/all/instance.rs | 4 +- tests/all/limits.rs | 14 +- tests/all/main.rs | 16 + tests/all/memory.rs | 4 +- tests/all/pooling_allocator.rs | 448 +++++- tests/all/wast.rs | 8 +- 32 files changed, 2619 insertions(+), 1489 deletions(-) create mode 100644 crates/runtime/src/instance/allocator/pooling/memory_pool.rs create mode 100644 crates/runtime/src/instance/allocator/pooling/simple.rs create mode 100644 crates/runtime/src/instance/allocator/pooling/stack_pool.rs create mode 100644 crates/runtime/src/instance/allocator/pooling/table_pool.rs diff --git a/RELEASES.md b/RELEASES.md index 3e0d3d10feaf..ef44e8f97791 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -8,6 +8,51 @@ Unreleased. ### Changed +* The pooling allocator was significantly refactored and the + `PoolingAllocationConfig` has some minor breaking API changes that reflect + those changes. + + Previously, the pooling allocator had `count` slots, and each slot had `N` + memories and `M` tables. Every allocated instance would reserve those `N` + memories and `M` tables regardless whether it actually needed them all or + not. This could lead to some waste and over-allocation when a module used less + memories and tables than the pooling allocator's configured maximums. + + After the refactors in this release, the pooling allocator doesn't have + one-size-fits-all slots anymore. Instead, memories and tables are in separate + pools that can be allocated from independently, and we allocate exactly as + many memories and tables as are necessary for the instance being allocated. + + To preserve your old configuration with the new methods you can do the following: + + ```rust + let mut config = PoolingAllocationConfig::default(); + + // If you used to have this old, no-longer-compiling configuration: + config.count(count); + config.instance_memories(n); + config.instance_tables(m); + + // You can use these equivalent settings for the new config methods: + config.total_core_instances(count); + config.total_stacks(count); // If using the `async` feature. + config.total_memories(count * n); + config.max_memories_per_module(n); + config.total_tables(count * m); + config.max_tables_per_module(m); + ``` + + There are additionally a variety of methods to limit the maximum amount of + resources a single core Wasm or component instance can take from the pool: + + * `PoolingAllocationConfig::max_memories_per_module` + * `PoolingAllocationConfig::max_tables_per_module` + * `PoolingAllocationConfig::max_memories_per_component` + * `PoolingAllocationConfig::max_tables_per_component` + * `PoolingAllocationConfig::max_core_instances_per_component` + + These methods do not affect the size of the pre-allocated pool. + * Options to the `wasmtime` CLI for Wasmtime itself must now come before the WebAssembly module. For example `wasmtime run foo.wasm --disable-cache` now must be specified as `wasmtime run --disable-cache foo.wasm`. Any diff --git a/benches/instantiation.rs b/benches/instantiation.rs index c1810ac166d5..e43bd650545e 100644 --- a/benches/instantiation.rs +++ b/benches/instantiation.rs @@ -219,7 +219,7 @@ fn strategies() -> impl Iterator { InstanceAllocationStrategy::OnDemand, InstanceAllocationStrategy::Pooling({ let mut config = PoolingAllocationConfig::default(); - config.instance_memory_pages(10_000); + config.memory_pages(10_000); config }), ] diff --git a/crates/fuzzing/src/generators/config.rs b/crates/fuzzing/src/generators/config.rs index 07daebc8083d..08c1967145b6 100644 --- a/crates/fuzzing/src/generators/config.rs +++ b/crates/fuzzing/src/generators/config.rs @@ -73,13 +73,13 @@ impl Config { // If using the pooling allocator, update the instance limits too if let InstanceAllocationStrategy::Pooling(pooling) = &mut self.wasmtime.strategy { // One single-page memory - pooling.instance_memories = config.max_memories as u32; - pooling.instance_memory_pages = 10; + pooling.total_memories = config.max_memories as u32; + pooling.memory_pages = 10; - pooling.instance_tables = config.max_tables as u32; - pooling.instance_table_elements = 1_000; + pooling.total_tables = config.max_tables as u32; + pooling.table_elements = 1_000; - pooling.instance_size = 1_000_000; + pooling.core_instance_size = 1_000_000; } } @@ -126,12 +126,12 @@ impl Config { if let InstanceAllocationStrategy::Pooling(pooling) = &self.wasmtime.strategy { // Check to see if any item limit is less than the required // threshold to execute the spec tests. - if pooling.instance_memories < 1 - || pooling.instance_tables < 5 - || pooling.instance_table_elements < 1_000 - || pooling.instance_memory_pages < 900 - || pooling.instance_count < 500 - || pooling.instance_size < 64 * 1024 + if pooling.total_memories < 1 + || pooling.total_tables < 5 + || pooling.table_elements < 1_000 + || pooling.memory_pages < 900 + || pooling.total_core_instances < 500 + || pooling.core_instance_size < 64 * 1024 { return false; } @@ -333,23 +333,23 @@ impl<'a> Arbitrary<'a> for Config { // Ensure the pooling allocator can support the maximal size of // memory, picking the smaller of the two to win. - if cfg.max_memory_pages < pooling.instance_memory_pages { - pooling.instance_memory_pages = cfg.max_memory_pages; + if cfg.max_memory_pages < pooling.memory_pages { + pooling.memory_pages = cfg.max_memory_pages; } else { - cfg.max_memory_pages = pooling.instance_memory_pages; + cfg.max_memory_pages = pooling.memory_pages; } // If traps are disallowed then memories must have at least one page // of memory so if we still are only allowing 0 pages of memory then // increase that to one here. if cfg.disallow_traps { - if pooling.instance_memory_pages == 0 { - pooling.instance_memory_pages = 1; + if pooling.memory_pages == 0 { + pooling.memory_pages = 1; cfg.max_memory_pages = 1; } // .. additionally update tables - if pooling.instance_table_elements == 0 { - pooling.instance_table_elements = 1; + if pooling.table_elements == 0 { + pooling.table_elements = 1; } } @@ -366,8 +366,8 @@ impl<'a> Arbitrary<'a> for Config { // Force this pooling allocator to always be able to accommodate the // module that may be generated. - pooling.instance_memories = cfg.max_memories as u32; - pooling.instance_tables = cfg.max_tables as u32; + pooling.total_memories = cfg.max_memories as u32; + pooling.total_tables = cfg.max_tables as u32; } Ok(config) diff --git a/crates/fuzzing/src/generators/pooling_config.rs b/crates/fuzzing/src/generators/pooling_config.rs index f670746dcebc..4eead14cad56 100644 --- a/crates/fuzzing/src/generators/pooling_config.rs +++ b/crates/fuzzing/src/generators/pooling_config.rs @@ -6,17 +6,30 @@ use arbitrary::{Arbitrary, Unstructured}; #[derive(Debug, Clone, Eq, PartialEq, Hash)] #[allow(missing_docs)] pub struct PoolingAllocationConfig { + pub total_component_instances: u32, + pub total_core_instances: u32, + pub total_memories: u32, + pub total_tables: u32, + pub total_stacks: u32, + + pub memory_pages: u64, + pub table_elements: u32, + + pub component_instance_size: usize, + pub max_memories_per_component: u32, + pub max_tables_per_component: u32, + + pub core_instance_size: usize, + pub max_memories_per_module: u32, + pub max_tables_per_module: u32, + + pub table_keep_resident: usize, + pub linear_memory_keep_resident: usize, + pub max_unused_warm_slots: u32, - pub instance_count: u32, - pub instance_memories: u32, - pub instance_tables: u32, - pub instance_memory_pages: u64, - pub instance_table_elements: u32, - pub instance_size: usize, + pub async_stack_zeroing: bool, pub async_stack_keep_resident: usize, - pub linear_memory_keep_resident: usize, - pub table_keep_resident: usize, } impl PoolingAllocationConfig { @@ -24,17 +37,31 @@ impl PoolingAllocationConfig { pub fn to_wasmtime(&self) -> wasmtime::PoolingAllocationConfig { let mut cfg = wasmtime::PoolingAllocationConfig::default(); - cfg.max_unused_warm_slots(self.max_unused_warm_slots) - .instance_count(self.instance_count) - .instance_memories(self.instance_memories) - .instance_tables(self.instance_tables) - .instance_memory_pages(self.instance_memory_pages) - .instance_table_elements(self.instance_table_elements) - .instance_size(self.instance_size) - .async_stack_zeroing(self.async_stack_zeroing) - .async_stack_keep_resident(self.async_stack_keep_resident) - .linear_memory_keep_resident(self.linear_memory_keep_resident) - .table_keep_resident(self.table_keep_resident); + cfg.total_component_instances(self.total_component_instances); + cfg.total_core_instances(self.total_core_instances); + cfg.total_memories(self.total_memories); + cfg.total_tables(self.total_tables); + cfg.total_stacks(self.total_stacks); + + cfg.memory_pages(self.memory_pages); + cfg.table_elements(self.table_elements); + + cfg.component_instance_size(self.component_instance_size); + cfg.max_memories_per_component(self.max_memories_per_component); + cfg.max_tables_per_component(self.max_tables_per_component); + + cfg.core_instance_size(self.core_instance_size); + cfg.max_memories_per_module(self.max_memories_per_module); + cfg.max_tables_per_module(self.max_tables_per_module); + + cfg.table_keep_resident(self.table_keep_resident); + cfg.linear_memory_keep_resident(self.linear_memory_keep_resident); + + cfg.max_unused_warm_slots(self.max_unused_warm_slots); + + cfg.async_stack_zeroing(self.async_stack_zeroing); + cfg.async_stack_keep_resident(self.async_stack_keep_resident); + cfg } } @@ -42,26 +69,41 @@ impl PoolingAllocationConfig { impl<'a> Arbitrary<'a> for PoolingAllocationConfig { fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result { const MAX_COUNT: u32 = 100; - const MAX_TABLES: u32 = 10; - const MAX_MEMORIES: u32 = 10; + const MAX_TABLES: u32 = 100; + const MAX_MEMORIES: u32 = 100; const MAX_ELEMENTS: u32 = 1000; const MAX_MEMORY_PAGES: u64 = 160; // 10 MiB const MAX_SIZE: usize = 1 << 20; // 1 MiB + const MAX_INSTANCE_MEMORIES: u32 = 10; + const MAX_INSTANCE_TABLES: u32 = 10; - let instance_count = u.int_in_range(1..=MAX_COUNT)?; + let total_memories = u.int_in_range(1..=MAX_MEMORIES)?; Ok(Self { - max_unused_warm_slots: u.int_in_range(0..=instance_count + 10)?, - instance_tables: u.int_in_range(0..=MAX_TABLES)?, - instance_memories: u.int_in_range(0..=MAX_MEMORIES)?, - instance_table_elements: u.int_in_range(0..=MAX_ELEMENTS)?, - instance_memory_pages: u.int_in_range(0..=MAX_MEMORY_PAGES)?, - instance_count, - instance_size: u.int_in_range(0..=MAX_SIZE)?, + total_component_instances: u.int_in_range(1..=MAX_COUNT)?, + total_core_instances: u.int_in_range(1..=MAX_COUNT)?, + total_memories, + total_tables: u.int_in_range(1..=MAX_TABLES)?, + total_stacks: u.int_in_range(1..=MAX_COUNT)?, + + memory_pages: u.int_in_range(0..=MAX_MEMORY_PAGES)?, + table_elements: u.int_in_range(0..=MAX_ELEMENTS)?, + + component_instance_size: u.int_in_range(0..=MAX_SIZE)?, + max_memories_per_component: u.int_in_range(1..=MAX_INSTANCE_MEMORIES)?, + max_tables_per_component: u.int_in_range(1..=MAX_INSTANCE_TABLES)?, + + core_instance_size: u.int_in_range(0..=MAX_SIZE)?, + max_memories_per_module: u.int_in_range(1..=MAX_INSTANCE_MEMORIES)?, + max_tables_per_module: u.int_in_range(1..=MAX_INSTANCE_TABLES)?, + + table_keep_resident: u.int_in_range(0..=1 << 20)?, + linear_memory_keep_resident: u.int_in_range(0..=1 << 20)?, + + max_unused_warm_slots: u.int_in_range(0..=total_memories + 10)?, + async_stack_zeroing: u.arbitrary()?, async_stack_keep_resident: u.int_in_range(0..=1 << 20)?, - linear_memory_keep_resident: u.int_in_range(0..=1 << 20)?, - table_keep_resident: u.int_in_range(0..=1 << 20)?, }) } } diff --git a/crates/jit/src/instantiate.rs b/crates/jit/src/instantiate.rs index c108f9b97f8d..3330245f6170 100644 --- a/crates/jit/src/instantiate.rs +++ b/crates/jit/src/instantiate.rs @@ -56,7 +56,7 @@ impl CompiledFunctionInfo { #[derive(Serialize, Deserialize)] pub struct CompiledModuleInfo { /// Type information about the compiled WebAssembly module. - module: Module, + pub module: Module, /// Metadata about each compiled function. funcs: PrimaryMap, diff --git a/crates/runtime/src/instance.rs b/crates/runtime/src/instance.rs index 7a4c48ffa69e..d6db99df9e70 100644 --- a/crates/runtime/src/instance.rs +++ b/crates/runtime/src/instance.rs @@ -67,13 +67,21 @@ pub struct Instance { /// /// This is where all runtime information about defined linear memories in /// this module lives. - memories: PrimaryMap, + /// + /// The `MemoryAllocationIndex` was given from our `InstanceAllocator` and + /// must be given back to the instance allocator when deallocating each + /// memory. + memories: PrimaryMap, /// WebAssembly table data. /// /// Like memories, this is only for defined tables in the module and /// contains all of their runtime state. - tables: PrimaryMap, + /// + /// The `TableAllocationIndex` was given from our `InstanceAllocator` and + /// must be given back to the instance allocator when deallocating each + /// table. + tables: PrimaryMap, /// Stores the dropped passive element segments in this instantiation by index. /// If the index is present in the set, the segment has been dropped. @@ -89,13 +97,6 @@ pub struct Instance { /// allocation, but some host-defined objects will store their state here. host_state: Box, - /// Instance of this instance within its `InstanceAllocator` trait - /// implementation. - /// - /// This is always 0 for the on-demand instance allocator and it's the - /// index of the slot in the pooling allocator. - index: usize, - /// A pointer to the `vmctx` field at the end of the `Instance`. /// /// If you're looking at this a reasonable question would be "why do we need @@ -160,9 +161,8 @@ impl Instance { /// allocation was `alloc_size` in bytes. unsafe fn new( req: InstanceAllocationRequest, - index: usize, - memories: PrimaryMap, - tables: PrimaryMap, + memories: PrimaryMap, + tables: PrimaryMap, memory_plans: &PrimaryMap, ) -> InstanceHandle { // The allocation must be *at least* the size required of `Instance`. @@ -184,7 +184,6 @@ impl Instance { ptr, Instance { runtime_info: req.runtime_info.clone(), - index, memories, tables, dropped_elements, @@ -567,7 +566,7 @@ impl Instance { delta: u64, ) -> Result, Error> { let store = unsafe { &mut *self.store() }; - let memory = &mut self.memories[idx]; + let memory = &mut self.memories[idx].1; let result = unsafe { memory.grow(delta, Some(store)) }; @@ -608,16 +607,17 @@ impl Instance { init_value: TableElement, ) -> Result, Error> { let store = unsafe { &mut *self.store() }; - let table = self + let table = &mut self .tables .get_mut(table_index) - .unwrap_or_else(|| panic!("no table for index {}", table_index.index())); + .unwrap_or_else(|| panic!("no table for index {}", table_index.index())) + .1; let result = unsafe { table.grow(delta, init_value, store) }; // Keep the `VMContext` pointers used by compiled Wasm code up to // date. - let element = self.tables[table_index].vmtable(); + let element = self.tables[table_index].1.vmtable(); self.set_table(table_index, element); result @@ -806,7 +806,7 @@ impl Instance { /// Get a locally-defined memory. pub fn get_defined_memory(&mut self, index: DefinedMemoryIndex) -> *mut Memory { - ptr::addr_of_mut!(self.memories[index]) + ptr::addr_of_mut!(self.memories[index].1) } /// Do a `memory.copy` @@ -975,11 +975,11 @@ impl Instance { idx: DefinedTableIndex, range: impl Iterator, ) -> *mut Table { - let elt_ty = self.tables[idx].element_type(); + let elt_ty = self.tables[idx].1.element_type(); if elt_ty == TableElementType::Func { for i in range { - let value = match self.tables[idx].get(i) { + let value = match self.tables[idx].1.get(i) { Some(value) => value, None => { // Out-of-bounds; caller will handle by likely @@ -1010,25 +1010,26 @@ impl Instance { .and_then(|func_index| self.get_func_ref(func_index)) .unwrap_or(std::ptr::null_mut()); self.tables[idx] + .1 .set(i, TableElement::FuncRef(func_ref)) .expect("Table type should match and index should be in-bounds"); } } - ptr::addr_of_mut!(self.tables[idx]) + ptr::addr_of_mut!(self.tables[idx].1) } /// Get a table by index regardless of whether it is locally-defined or an /// imported, foreign table. pub(crate) fn get_table(&mut self, table_index: TableIndex) -> *mut Table { self.with_defined_table_index_and_instance(table_index, |idx, instance| { - ptr::addr_of_mut!(instance.tables[idx]) + ptr::addr_of_mut!(instance.tables[idx].1) }) } /// Get a locally-defined table. pub(crate) fn get_defined_table(&mut self, index: DefinedTableIndex) -> *mut Table { - ptr::addr_of_mut!(self.tables[index]) + ptr::addr_of_mut!(self.tables[index].1) } pub(crate) fn with_defined_table_index_and_instance( @@ -1110,7 +1111,7 @@ impl Instance { // Initialize the defined tables let mut ptr = self.vmctx_plus_offset_mut(offsets.vmctx_tables_begin()); for i in 0..module.table_plans.len() - module.num_imported_tables { - ptr::write(ptr, self.tables[DefinedTableIndex::new(i)].vmtable()); + ptr::write(ptr, self.tables[DefinedTableIndex::new(i)].1.vmtable()); ptr = ptr.add(1); } @@ -1126,12 +1127,13 @@ impl Instance { let memory_index = module.memory_index(defined_memory_index); if module.memory_plans[memory_index].memory.shared { let def_ptr = self.memories[defined_memory_index] + .1 .as_shared_memory() .unwrap() .vmmemory_ptr(); ptr::write(ptr, def_ptr.cast_mut()); } else { - ptr::write(owned_ptr, self.memories[defined_memory_index].vmmemory()); + ptr::write(owned_ptr, self.memories[defined_memory_index].1.vmmemory()); ptr::write(ptr, owned_ptr); owned_ptr = owned_ptr.add(1); } @@ -1198,7 +1200,7 @@ impl Instance { fn wasm_fault(&self, addr: usize) -> Option { let mut fault = None; - for (_, memory) in self.memories.iter() { + for (_, (_, memory)) in self.memories.iter() { let accessible = memory.wasm_accessible(); if accessible.start <= addr && addr < accessible.end { // All linear memories should be disjoint so assert that no diff --git a/crates/runtime/src/instance/allocator.rs b/crates/runtime/src/instance/allocator.rs index 4ed72d6964b1..5388efb4ed2e 100644 --- a/crates/runtime/src/instance/allocator.rs +++ b/crates/runtime/src/instance/allocator.rs @@ -4,15 +4,17 @@ use crate::memory::Memory; use crate::table::Table; use crate::{CompiledModuleId, ModuleRuntimeInfo, Store}; use anyhow::{anyhow, bail, Result}; -use std::alloc; -use std::any::Any; -use std::convert::TryFrom; -use std::ptr; -use std::sync::Arc; +use std::{alloc, any::Any, mem, ptr, sync::Arc}; use wasmtime_environ::{ DefinedMemoryIndex, DefinedTableIndex, HostPtr, InitMemory, MemoryInitialization, - MemoryInitializer, Module, PrimaryMap, TableInitialValue, TableSegment, Trap, VMOffsets, - WasmType, WASM_PAGE_SIZE, + MemoryInitializer, MemoryPlan, Module, PrimaryMap, TableInitialValue, TablePlan, TableSegment, + Trap, VMOffsets, WasmType, WASM_PAGE_SIZE, +}; + +#[cfg(feature = "component-model")] +use wasmtime_environ::{ + component::{Component, VMComponentOffsets}, + StaticModuleIndex, }; mod on_demand; @@ -64,20 +66,25 @@ pub struct InstanceAllocationRequest<'a> { /// itself, because several use-sites require a split mut borrow on the /// InstanceAllocationRequest. pub struct StorePtr(Option<*mut dyn Store>); + impl StorePtr { /// A pointer to no Store. pub fn empty() -> Self { Self(None) } + /// A pointer to a Store. pub fn new(ptr: *mut dyn Store) -> Self { Self(Some(ptr)) } + /// The raw contents of this struct pub fn as_raw(&self) -> Option<*mut dyn Store> { self.0.clone() } + /// Use the StorePtr as a mut ref to the Store. + /// /// Safety: must not be used outside the original lifetime of the borrow. pub(crate) unsafe fn get(&mut self) -> Option<&mut dyn Store> { match self.0 { @@ -87,16 +94,187 @@ impl StorePtr { } } -/// Represents a runtime instance allocator. +/// The index of a memory allocation within an `InstanceAllocator`. +#[derive(Clone, Copy, Debug, Eq, PartialEq, PartialOrd, Ord)] +pub struct MemoryAllocationIndex(u32); + +impl Default for MemoryAllocationIndex { + fn default() -> Self { + // A default `MemoryAllocationIndex` that can be used with + // `InstanceAllocator`s that don't actually need indices. + MemoryAllocationIndex(u32::MAX) + } +} + +impl MemoryAllocationIndex { + /// Get the underlying index of this `MemoryAllocationIndex`. + pub fn index(&self) -> usize { + self.0 as usize + } +} + +/// The index of a table allocation within an `InstanceAllocator`. +#[derive(Clone, Copy, Debug, Eq, PartialEq, PartialOrd, Ord)] +pub struct TableAllocationIndex(u32); + +impl Default for TableAllocationIndex { + fn default() -> Self { + // A default `TableAllocationIndex` that can be used with + // `InstanceAllocator`s that don't actually need indices. + TableAllocationIndex(u32::MAX) + } +} + +impl TableAllocationIndex { + /// Get the underlying index of this `TableAllocationIndex`. + pub fn index(&self) -> usize { + self.0 as usize + } +} + +/// Trait that represents the hooks needed to implement an instance allocator. +/// +/// Implement this trait when implementing new instance allocators, but don't +/// use this trait when you need an instance allocator. Instead use the +/// `InstanceAllocator` trait for that, which has additional helper methods and +/// a blanket implementation for all types that implement this trait. /// /// # Safety /// -/// This trait is unsafe as it requires knowledge of Wasmtime's runtime internals to implement correctly. -pub unsafe trait InstanceAllocator { - /// Validates that a module is supported by the allocator. - fn validate(&self, module: &Module, offsets: &VMOffsets) -> Result<()> { - let _ = (module, offsets); - Ok(()) +/// This trait is unsafe as it requires knowledge of Wasmtime's runtime +/// internals to implement correctly. +pub unsafe trait InstanceAllocatorImpl { + /// Validate whether a component (including all of its contained core + /// modules) is allocatable by this instance allocator. + #[cfg(feature = "component-model")] + fn validate_component_impl<'a>( + &self, + component: &Component, + offsets: &VMComponentOffsets, + get_module: &'a dyn Fn(StaticModuleIndex) -> &'a Module, + ) -> Result<()>; + + /// Validate whether a module is allocatable by this instance allocator. + fn validate_module_impl(&self, module: &Module, offsets: &VMOffsets) -> Result<()>; + + /// Increment the count of concurrent component instances that are currently + /// allocated, if applicable. + /// + /// Not all instance allocators will have limits for the maximum number of + /// concurrent component instances that can be live at the same time, and + /// these allocators may implement this method with a no-op. + fn increment_component_instance_count(&self) -> Result<()>; + + /// The dual of `increment_component_instance_count`. + fn decrement_component_instance_count(&self); + + /// Increment the count of concurrent core module instances that are + /// currently allocated, if applicable. + /// + /// Not all instance allocators will have limits for the maximum number of + /// concurrent core module instances that can be live at the same time, and + /// these allocators may implement this method with a no-op. + fn increment_core_instance_count(&self) -> Result<()>; + + /// The dual of `increment_core_instance_count`. + fn decrement_core_instance_count(&self); + + /// Allocate a memory for an instance. + /// + /// # Unsafety + /// + /// The memory and its associated module must have already been validated by + /// `Self::validate_module` and passed that validation. + unsafe fn allocate_memory( + &self, + request: &mut InstanceAllocationRequest, + memory_plan: &MemoryPlan, + memory_index: DefinedMemoryIndex, + ) -> Result<(MemoryAllocationIndex, Memory)>; + + /// Deallocate an instance's previously allocated memory. + /// + /// # Unsafety + /// + /// The memory must have previously been allocated by + /// `Self::allocate_memory`, be at the given index, and must currently be + /// allocated. It must never be used again. + unsafe fn deallocate_memory( + &self, + memory_index: DefinedMemoryIndex, + allocation_index: MemoryAllocationIndex, + memory: Memory, + ); + + /// Allocate a table for an instance. + /// + /// # Unsafety + /// + /// The table and its associated module must have already been validated by + /// `Self::validate_module` and passed that validation. + unsafe fn allocate_table( + &self, + req: &mut InstanceAllocationRequest, + table_plan: &TablePlan, + table_index: DefinedTableIndex, + ) -> Result<(TableAllocationIndex, Table)>; + + /// Deallocate an instance's previously allocated table. + /// + /// # Unsafety + /// + /// The table must have previously been allocated by `Self::allocate_table`, + /// be at the given index, and must currently be allocated. It must never be + /// used again. + unsafe fn deallocate_table( + &self, + table_index: DefinedTableIndex, + allocation_index: TableAllocationIndex, + table: Table, + ); + + /// Allocates a fiber stack for calling async functions on. + #[cfg(feature = "async")] + fn allocate_fiber_stack(&self) -> Result; + + /// Deallocates a fiber stack that was previously allocated with `allocate_fiber_stack`. + /// + /// # Safety + /// + /// The provided stack is required to have been allocated with `allocate_fiber_stack`. + #[cfg(feature = "async")] + unsafe fn deallocate_fiber_stack(&self, stack: &wasmtime_fiber::FiberStack); + + /// Purges all lingering resources related to `module` from within this + /// allocator. + /// + /// Primarily present for the pooling allocator to remove mappings of + /// this module from slots in linear memory. + fn purge_module(&self, module: CompiledModuleId); +} + +/// A thing that can allocate instances. +/// +/// Don't implement this trait directly, instead implement +/// `InstanceAllocatorImpl` and you'll get this trait for free via a blanket +/// impl. +pub trait InstanceAllocator: InstanceAllocatorImpl { + /// Validate whether a component (including all of its contained core + /// modules) is allocatable with this instance allocator. + #[cfg(feature = "component-model")] + fn validate_component<'a>( + &self, + component: &Component, + offsets: &VMComponentOffsets, + get_module: &'a dyn Fn(StaticModuleIndex) -> &'a Module, + ) -> Result<()> { + InstanceAllocatorImpl::validate_component_impl(self, component, offsets, get_module) + } + + /// Validate whether a core module is allocatable with this instance + /// allocator. + fn validate_module(&self, module: &Module, offsets: &VMOffsets) -> Result<()> { + InstanceAllocatorImpl::validate_module_impl(self, module, offsets) } /// Allocates a fresh `InstanceHandle` for the `req` given. @@ -107,28 +285,42 @@ pub unsafe trait InstanceAllocator { /// /// Note that the returned instance must still have `.initialize(..)` called /// on it to complete the instantiation process. - fn allocate(&self, mut req: InstanceAllocationRequest) -> Result { - let index = self.allocate_index(&req)?; - let module = req.runtime_info.module(); - let mut memories = - PrimaryMap::with_capacity(module.memory_plans.len() - module.num_imported_memories); - let mut tables = - PrimaryMap::with_capacity(module.table_plans.len() - module.num_imported_tables); + /// + /// # Unsafety + /// + /// The request's associated module, memories, tables, and vmctx must have + /// already have been validated by `Self::validate_module`. + unsafe fn allocate_module( + &self, + mut request: InstanceAllocationRequest, + ) -> Result { + let module = request.runtime_info.module(); + + #[cfg(debug_assertions)] + InstanceAllocatorImpl::validate_module_impl(self, module, request.runtime_info.offsets()) + .expect("module should have already been validated before allocation"); + + self.increment_core_instance_count()?; + + let num_defined_memories = module.memory_plans.len() - module.num_imported_memories; + let mut memories = PrimaryMap::with_capacity(num_defined_memories); + + let num_defined_tables = module.table_plans.len() - module.num_imported_tables; + let mut tables = PrimaryMap::with_capacity(num_defined_tables); let result = self - .allocate_memories(index, &mut req, &mut memories) - .and_then(|()| self.allocate_tables(index, &mut req, &mut tables)); + .allocate_memories(&mut request, &mut memories) + .and_then(|()| self.allocate_tables(&mut request, &mut tables)); if let Err(e) = result { - self.deallocate_memories(index, &mut memories); - self.deallocate_tables(index, &mut tables); - self.deallocate_index(index); + self.deallocate_memories(&mut memories); + self.deallocate_tables(&mut tables); + self.decrement_core_instance_count(); return Err(e); } unsafe { Ok(Instance::new( - req, - index, + request, memories, tables, &module.memory_plans, @@ -140,79 +332,120 @@ pub unsafe trait InstanceAllocator { /// /// This will null-out the pointer within `handle` and otherwise reclaim /// resources such as tables, memories, and the instance memory itself. - fn deallocate(&self, handle: &mut InstanceHandle) { - let index = handle.instance().index; - self.deallocate_memories(index, &mut handle.instance_mut().memories); - self.deallocate_tables(index, &mut handle.instance_mut().tables); - unsafe { - let layout = Instance::alloc_layout(handle.instance().offsets()); - let ptr = handle.instance.take().unwrap(); - ptr::drop_in_place(ptr.as_ptr()); - alloc::dealloc(ptr.as_ptr().cast(), layout); - } - self.deallocate_index(index); + /// + /// # Unsafety + /// + /// The instance must have previously been allocated by `Self::allocate`. + unsafe fn deallocate_module(&self, handle: &mut InstanceHandle) { + self.deallocate_memories(&mut handle.instance_mut().memories); + self.deallocate_tables(&mut handle.instance_mut().tables); + + let layout = Instance::alloc_layout(handle.instance().offsets()); + let ptr = handle.instance.take().unwrap(); + ptr::drop_in_place(ptr.as_ptr()); + alloc::dealloc(ptr.as_ptr().cast(), layout); + + self.decrement_core_instance_count(); } - /// Optionally allocates an allocator-defined index for the `req` provided. + /// Allocate the memories for the given instance allocation request, pushing + /// them into `memories`. + /// + /// # Unsafety /// - /// The return value here, if successful, is passed to the various methods - /// below for memory/table allocation/deallocation. - fn allocate_index(&self, req: &InstanceAllocationRequest) -> Result; + /// The request's associated module and memories must have previously been + /// validated by `Self::validate_module`. + unsafe fn allocate_memories( + &self, + request: &mut InstanceAllocationRequest, + memories: &mut PrimaryMap, + ) -> Result<()> { + let module = request.runtime_info.module(); + + #[cfg(debug_assertions)] + InstanceAllocatorImpl::validate_module_impl(self, module, request.runtime_info.offsets()) + .expect("module should have already been validated before allocation"); + + for (memory_index, memory_plan) in module + .memory_plans + .iter() + .skip(module.num_imported_memories) + { + let memory_index = module + .defined_memory_index(memory_index) + .expect("should be a defined memory since we skipped imported ones"); + + memories.push(self.allocate_memory(request, memory_plan, memory_index)?); + } - /// Deallocates indices allocated by `allocate_index`. - fn deallocate_index(&self, index: usize); + Ok(()) + } - /// Attempts to allocate all defined linear memories for a module. + /// Deallocate all the memories in the given primary map. /// - /// Pushes all memories for `req` onto the `mems` storage provided which is - /// already appropriately allocated to contain all memories. + /// # Unsafety /// - /// Note that this is allowed to fail. Failure can additionally happen after - /// some memories have already been successfully allocated. All memories - /// pushed onto `mem` are guaranteed to one day make their way to - /// `deallocate_memories`. - fn allocate_memories( + /// The memories must have previously been allocated by + /// `Self::allocate_memories`. + unsafe fn deallocate_memories( &self, - index: usize, - req: &mut InstanceAllocationRequest, - mems: &mut PrimaryMap, - ) -> Result<()>; - - /// Deallocates all memories provided, optionally reclaiming resources for - /// the pooling allocator for example. - fn deallocate_memories(&self, index: usize, mems: &mut PrimaryMap); + memories: &mut PrimaryMap, + ) { + for (memory_index, (allocation_index, memory)) in mem::take(memories) { + self.deallocate_memory(memory_index, allocation_index, memory); + } + } - /// Same as `allocate_memories`, but for tables. - fn allocate_tables( + /// Allocate tables for the given instance allocation request, pushing them + /// into `tables`. + /// + /// # Unsafety + /// + /// The request's associated module and tables must have previously been + /// validated by `Self::validate_module`. + unsafe fn allocate_tables( &self, - index: usize, - req: &mut InstanceAllocationRequest, - tables: &mut PrimaryMap, - ) -> Result<()>; + request: &mut InstanceAllocationRequest, + tables: &mut PrimaryMap, + ) -> Result<()> { + let module = request.runtime_info.module(); - /// Same as `deallocate_memories`, but for tables. - fn deallocate_tables(&self, index: usize, tables: &mut PrimaryMap); + #[cfg(debug_assertions)] + InstanceAllocatorImpl::validate_module_impl(self, module, request.runtime_info.offsets()) + .expect("module should have already been validated before allocation"); - /// Allocates a fiber stack for calling async functions on. - #[cfg(feature = "async")] - fn allocate_fiber_stack(&self) -> Result; + for (index, plan) in module.table_plans.iter().skip(module.num_imported_tables) { + let def_index = module + .defined_table_index(index) + .expect("should be a defined table since we skipped imported ones"); - /// Deallocates a fiber stack that was previously allocated with `allocate_fiber_stack`. - /// - /// # Safety - /// - /// The provided stack is required to have been allocated with `allocate_fiber_stack`. - #[cfg(feature = "async")] - unsafe fn deallocate_fiber_stack(&self, stack: &wasmtime_fiber::FiberStack); + tables.push(self.allocate_table(request, plan, def_index)?); + } - /// Purges all lingering resources related to `module` from within this - /// allocator. + Ok(()) + } + + /// Deallocate all the tables in the given primary map. /// - /// Primarily present for the pooling allocator to remove mappings of - /// this module from slots in linear memory. - fn purge_module(&self, module: CompiledModuleId); + /// # Unsafety + /// + /// The tables must have previously been allocated by + /// `Self::allocate_tables`. + unsafe fn deallocate_tables( + &self, + tables: &mut PrimaryMap, + ) { + for (table_index, (allocation_index, table)) in mem::take(tables) { + self.deallocate_table(table_index, allocation_index, table); + } + } } +// Every `InstanceAllocatorImpl` is an `InstanceAllocator` when used +// correctly. Also, no one is allowed to override this trait's methods, they +// must use the defaults. This blanket impl provides both of those things. +impl InstanceAllocator for T {} + fn get_table_init_start(init: &TableSegment, instance: &mut Instance) -> Result { match init.base { Some(base) => { @@ -364,7 +597,7 @@ fn initialize_memories(instance: &mut Instance, module: &Module) -> Result<()> { // pre-initializing it via mmap magic, then this initializer can be // skipped entirely. if let Some(memory_index) = module.defined_memory_index(memory_index) { - if !instance.memories[memory_index].needs_init() { + if !instance.memories[memory_index].1.needs_init() { return true; } } @@ -423,3 +656,14 @@ pub(super) fn initialize_instance( Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn allocator_traits_are_object_safe() { + fn _instance_allocator(_: &dyn InstanceAllocatorImpl) {} + fn _instance_allocator_ext(_: &dyn InstanceAllocator) {} + } +} diff --git a/crates/runtime/src/instance/allocator/on_demand.rs b/crates/runtime/src/instance/allocator/on_demand.rs index 171562f9d918..ad4d951bf57a 100644 --- a/crates/runtime/src/instance/allocator/on_demand.rs +++ b/crates/runtime/src/instance/allocator/on_demand.rs @@ -1,11 +1,21 @@ -use super::{InstanceAllocationRequest, InstanceAllocator}; +use super::{ + InstanceAllocationRequest, InstanceAllocatorImpl, MemoryAllocationIndex, TableAllocationIndex, +}; use crate::instance::RuntimeMemoryCreator; use crate::memory::{DefaultMemoryCreator, Memory}; use crate::table::Table; use crate::CompiledModuleId; use anyhow::Result; use std::sync::Arc; -use wasmtime_environ::{DefinedMemoryIndex, DefinedTableIndex, PrimaryMap}; +use wasmtime_environ::{ + DefinedMemoryIndex, DefinedTableIndex, HostPtr, MemoryPlan, Module, TablePlan, VMOffsets, +}; + +#[cfg(feature = "component-model")] +use wasmtime_environ::{ + component::{Component, VMComponentOffsets}, + StaticModuleIndex, +}; /// Represents the on-demand instance allocator. #[derive(Clone)] @@ -37,75 +47,92 @@ impl Default for OnDemandInstanceAllocator { } } -unsafe impl InstanceAllocator for OnDemandInstanceAllocator { - fn allocate_index(&self, _req: &InstanceAllocationRequest) -> Result { - Ok(0) +unsafe impl InstanceAllocatorImpl for OnDemandInstanceAllocator { + #[cfg(feature = "component-model")] + fn validate_component_impl<'a>( + &self, + _component: &Component, + _offsets: &VMComponentOffsets, + _get_module: &'a dyn Fn(StaticModuleIndex) -> &'a Module, + ) -> Result<()> { + Ok(()) } - fn deallocate_index(&self, index: usize) { - assert_eq!(index, 0); + fn validate_module_impl(&self, _module: &Module, _offsets: &VMOffsets) -> Result<()> { + Ok(()) } - fn allocate_memories( + fn increment_component_instance_count(&self) -> Result<()> { + Ok(()) + } + + fn decrement_component_instance_count(&self) {} + + fn increment_core_instance_count(&self) -> Result<()> { + Ok(()) + } + + fn decrement_core_instance_count(&self) {} + + unsafe fn allocate_memory( &self, - _index: usize, - req: &mut InstanceAllocationRequest, - memories: &mut PrimaryMap, - ) -> Result<()> { - let module = req.runtime_info.module(); + request: &mut InstanceAllocationRequest, + memory_plan: &MemoryPlan, + memory_index: DefinedMemoryIndex, + ) -> Result<(MemoryAllocationIndex, Memory)> { let creator = self .mem_creator .as_deref() .unwrap_or_else(|| &DefaultMemoryCreator); - let num_imports = module.num_imported_memories; - for (memory_idx, plan) in module.memory_plans.iter().skip(num_imports) { - let defined_memory_idx = module - .defined_memory_index(memory_idx) - .expect("Skipped imports, should never be None"); - let image = req.runtime_info.memory_image(defined_memory_idx)?; - - memories.push(Memory::new_dynamic( - plan, - creator, - unsafe { - req.store - .get() - .expect("if module has memory plans, store is not empty") - }, - image, - )?); - } - Ok(()) + let image = request.runtime_info.memory_image(memory_index)?; + let allocation_index = MemoryAllocationIndex::default(); + let memory = Memory::new_dynamic( + memory_plan, + creator, + request + .store + .get() + .expect("if module has memory plans, store is not empty"), + image, + )?; + Ok((allocation_index, memory)) } - fn deallocate_memories( + unsafe fn deallocate_memory( &self, - _index: usize, - _mems: &mut PrimaryMap, + _memory_index: DefinedMemoryIndex, + allocation_index: MemoryAllocationIndex, + _memory: Memory, ) { - // normal destructors do cleanup here + debug_assert_eq!(allocation_index, MemoryAllocationIndex::default()); + // Normal destructors do all the necessary clean up. } - fn allocate_tables( + unsafe fn allocate_table( &self, - _index: usize, - req: &mut InstanceAllocationRequest, - tables: &mut PrimaryMap, - ) -> Result<()> { - let module = req.runtime_info.module(); - let num_imports = module.num_imported_tables; - for (_, table) in module.table_plans.iter().skip(num_imports) { - tables.push(Table::new_dynamic(table, unsafe { - req.store - .get() - .expect("if module has table plans, store is not empty") - })?); - } - Ok(()) + request: &mut InstanceAllocationRequest, + table_plan: &TablePlan, + _table_index: DefinedTableIndex, + ) -> Result<(TableAllocationIndex, Table)> { + let allocation_index = TableAllocationIndex::default(); + let table = Table::new_dynamic( + table_plan, + request + .store + .get() + .expect("if module has table plans, store is not empty"), + )?; + Ok((allocation_index, table)) } - fn deallocate_tables(&self, _index: usize, _tables: &mut PrimaryMap) { - // normal destructors do cleanup here + unsafe fn deallocate_table( + &self, + _table_index: DefinedTableIndex, + allocation_index: TableAllocationIndex, + _table: Table, + ) { + debug_assert_eq!(allocation_index, TableAllocationIndex::default()); + // Normal destructors do all the necessary clean up. } #[cfg(feature = "async")] diff --git a/crates/runtime/src/instance/allocator/pooling.rs b/crates/runtime/src/instance/allocator/pooling.rs index 2ee511c6e76b..41a5579fc384 100644 --- a/crates/runtime/src/instance/allocator/pooling.rs +++ b/crates/runtime/src/instance/allocator/pooling.rs @@ -7,21 +7,12 @@ //! Using the pooling instance allocator can speed up module instantiation //! when modules can be constrained based on configurable limits. -use super::{InstanceAllocationRequest, InstanceAllocator}; -use crate::{instance::Instance, Memory, Mmap, Table}; -use crate::{CompiledModuleId, MemoryImageSlot}; -use anyhow::{anyhow, bail, Context, Result}; -use libc::c_void; -use std::convert::TryFrom; -use std::mem; -use std::sync::Mutex; -use wasmtime_environ::{ - DefinedMemoryIndex, DefinedTableIndex, HostPtr, MemoryStyle, Module, PrimaryMap, Tunables, - VMOffsets, WASM_PAGE_SIZE, -}; - mod index_allocator; -use index_allocator::{ModuleAffinityIndexAllocator, SlotId}; +mod memory_pool; +mod table_pool; + +#[cfg(all(feature = "async", unix, not(miri)))] +mod stack_pool; cfg_if::cfg_if! { if #[cfg(windows)] { @@ -33,10 +24,30 @@ cfg_if::cfg_if! { } } -use imp::{commit_table_pages, decommit_table_pages}; +use super::{ + InstanceAllocationRequest, InstanceAllocatorImpl, MemoryAllocationIndex, TableAllocationIndex, +}; +use crate::{instance::Instance, CompiledModuleId, Memory, Table}; +use anyhow::{bail, Result}; +use memory_pool::MemoryPool; +use std::{ + mem, + sync::atomic::{AtomicU64, Ordering}, +}; +use table_pool::TablePool; +use wasmtime_environ::{ + DefinedMemoryIndex, DefinedTableIndex, HostPtr, MemoryPlan, Module, TablePlan, Tunables, + VMOffsets, +}; #[cfg(all(feature = "async", unix, not(miri)))] -use imp::{commit_stack_pages, reset_stack_pages_to_zero}; +use stack_pool::StackPool; + +#[cfg(feature = "component-model")] +use wasmtime_environ::{ + component::{Component, VMComponentOffsets}, + StaticModuleIndex, +}; fn round_up_to_pow2(n: usize, to: usize) -> usize { debug_assert!(to > 0); @@ -49,22 +60,52 @@ fn round_up_to_pow2(n: usize, to: usize) -> usize { /// More docs on this can be found at `wasmtime::PoolingAllocationConfig`. #[derive(Debug, Copy, Clone)] pub struct InstanceLimits { - /// Maximum instances to support - pub count: u32, + /// The maximum number of component instances that may be allocated + /// concurrently. + pub total_component_instances: u32, + + /// The maximum size of a component's `VMComponentContext`, not including + /// any of its inner core modules' `VMContext` sizes. + pub component_instance_size: usize, + + /// The maximum number of core module instances that may be allocated + /// concurrently. + pub total_core_instances: u32, + + /// The maximum number of core module instances that a single component may + /// transitively contain. + pub max_core_instances_per_component: u32, + + /// The maximum number of Wasm linear memories that a component may + /// transitively contain. + pub max_memories_per_component: u32, + + /// The maximum number of tables that a component may transitively contain. + pub max_tables_per_component: u32, + + /// The total number of linear memories in the pool, across all instances. + pub total_memories: u32, + + /// The total number of tables in the pool, across all instances. + pub total_tables: u32, + + /// The total number of async stacks in the pool, across all instances. + #[cfg(feature = "async")] + pub total_stacks: u32, - /// Maximum size of instance VMContext - pub size: usize, + /// Maximum size of a core instance's `VMContext`. + pub core_instance_size: usize, - /// Maximum number of tables per instance - pub tables: u32, + /// Maximum number of tables per instance. + pub max_tables_per_module: u32, - /// Maximum number of table elements per table + /// Maximum number of table elements per table. pub table_elements: u32, - /// Maximum number of linear memories per instance - pub memories: u32, + /// Maximum number of linear memories per instance. + pub max_memories_per_module: u32, - /// Maximum number of wasm pages for each linear memory. + /// Maximum number of Wasm pages for each linear memory. pub memory_pages: u64, } @@ -73,451 +114,25 @@ impl Default for InstanceLimits { // See doc comments for `wasmtime::PoolingAllocationConfig` for these // default values Self { - count: 1000, - size: 1 << 20, // 1 MB - tables: 1, + total_component_instances: 1000, + component_instance_size: 1 << 20, // 1 MiB + total_core_instances: 1000, + max_core_instances_per_component: 20, + max_memories_per_component: 20, + max_tables_per_component: 20, + total_memories: 1000, + total_tables: 1000, + #[cfg(feature = "async")] + total_stacks: 1000, + core_instance_size: 1 << 20, // 1 MiB + max_tables_per_module: 1, table_elements: 10_000, - memories: 1, + max_memories_per_module: 1, memory_pages: 160, } } } -/// Represents a pool of WebAssembly linear memories. -/// -/// A linear memory is divided into accessible pages and guard pages. -/// -/// Each instance index into the pool returns an iterator over the base -/// addresses of the instance's linear memories. -/// -/// A diagram for this struct's fields is: -/// -/// ```ignore -/// memory_size -/// / -/// max_accessible / memory_and_guard_size -/// | / | -/// <--+---> / <-----------+----------> -/// <--------+-> -/// -/// +-----------+--------+---+-----------+ +--------+---+-----------+ -/// | PROT_NONE | | PROT_NONE | ... | | PROT_NONE | -/// +-----------+--------+---+-----------+ +--------+---+-----------+ -/// | |<------------------+----------------------------------> -/// \ | \ -/// mapping | `max_instances * max_memories` memories -/// / -/// initial_memory_offset -/// ``` -#[derive(Debug)] -struct MemoryPool { - mapping: Mmap, - // If using a copy-on-write allocation scheme, the slot management. We - // dynamically transfer ownership of a slot to a Memory when in - // use. - image_slots: Vec>>, - // The size, in bytes, of each linear memory's reservation, not including - // any guard region. - memory_size: usize, - // The size, in bytes, of each linear memory's reservation plus the trailing - // guard region allocated for it. - memory_and_guard_size: usize, - // The maximum size that can become accessible, in bytes, of each linear - // memory. Guaranteed to be a whole number of wasm pages. - max_accessible: usize, - // The size, in bytes, of the offset to the first linear memory in this - // pool. This is here to help account for the first region of guard pages, - // if desired, before the first linear memory. - initial_memory_offset: usize, - max_memories: usize, - max_instances: usize, -} - -impl MemoryPool { - fn new(instance_limits: &InstanceLimits, tunables: &Tunables) -> Result { - // The maximum module memory page count cannot exceed 65536 pages - if instance_limits.memory_pages > 0x10000 { - bail!( - "module memory page limit of {} exceeds the maximum of 65536", - instance_limits.memory_pages - ); - } - - // Interpret the larger of the maximal size of memory or the static - // memory bound as the size of the virtual address space reservation for - // memory itself. Typically `static_memory_bound` is 4G which helps - // elide most bounds checks in wasm. If `memory_pages` is larger, - // though, then this is a non-moving pooling allocator so create larger - // reservations for account for that. - let memory_size = instance_limits - .memory_pages - .max(tunables.static_memory_bound) - * u64::from(WASM_PAGE_SIZE); - - let memory_and_guard_size = - usize::try_from(memory_size + tunables.static_memory_offset_guard_size) - .map_err(|_| anyhow!("memory reservation size exceeds addressable memory"))?; - - assert!( - memory_and_guard_size % crate::page_size() == 0, - "memory size {} is not a multiple of system page size", - memory_and_guard_size - ); - - let max_instances = instance_limits.count as usize; - let max_memories = instance_limits.memories as usize; - let initial_memory_offset = if tunables.guard_before_linear_memory { - usize::try_from(tunables.static_memory_offset_guard_size).unwrap() - } else { - 0 - }; - - // The entire allocation here is the size of each memory times the - // max memories per instance times the number of instances allowed in - // this pool, plus guard regions. - // - // Note, though, that guard regions are required to be after each linear - // memory. If the `guard_before_linear_memory` setting is specified, - // then due to the contiguous layout of linear memories the guard pages - // after one memory are also guard pages preceding the next linear - // memory. This means that we only need to handle pre-guard-page sizes - // specially for the first linear memory, hence the - // `initial_memory_offset` variable here. If guards aren't specified - // before linear memories this is set to `0`, otherwise it's set to - // the same size as guard regions for other memories. - let allocation_size = memory_and_guard_size - .checked_mul(max_memories) - .and_then(|c| c.checked_mul(max_instances)) - .and_then(|c| c.checked_add(initial_memory_offset)) - .ok_or_else(|| { - anyhow!("total size of memory reservation exceeds addressable memory") - })?; - - // Create a completely inaccessible region to start - let mapping = Mmap::accessible_reserved(0, allocation_size) - .context("failed to create memory pool mapping")?; - - let num_image_slots = max_instances * max_memories; - let image_slots: Vec<_> = std::iter::repeat_with(|| Mutex::new(None)) - .take(num_image_slots) - .collect(); - - let pool = Self { - mapping, - image_slots, - memory_size: memory_size.try_into().unwrap(), - memory_and_guard_size, - initial_memory_offset, - max_memories, - max_instances, - max_accessible: (instance_limits.memory_pages as usize) * (WASM_PAGE_SIZE as usize), - }; - - Ok(pool) - } - - fn get_base(&self, instance_index: usize, memory_index: DefinedMemoryIndex) -> *mut u8 { - assert!(instance_index < self.max_instances); - let memory_index = memory_index.as_u32() as usize; - assert!(memory_index < self.max_memories); - let idx = instance_index * self.max_memories + memory_index; - let offset = self.initial_memory_offset + idx * self.memory_and_guard_size; - unsafe { self.mapping.as_ptr().offset(offset as isize).cast_mut() } - } - - #[cfg(test)] - fn get<'a>(&'a self, instance_index: usize) -> impl Iterator + 'a { - (0..self.max_memories) - .map(move |i| self.get_base(instance_index, DefinedMemoryIndex::from_u32(i as u32))) - } - - /// Take ownership of the given image slot. Must be returned via - /// `return_memory_image_slot` when the instance is done using it. - fn take_memory_image_slot( - &self, - instance_index: usize, - memory_index: DefinedMemoryIndex, - ) -> MemoryImageSlot { - let idx = instance_index * self.max_memories + (memory_index.as_u32() as usize); - let maybe_slot = self.image_slots[idx].lock().unwrap().take(); - - maybe_slot.unwrap_or_else(|| { - MemoryImageSlot::create( - self.get_base(instance_index, memory_index) as *mut c_void, - 0, - self.max_accessible, - ) - }) - } - - /// Return ownership of the given image slot. - fn return_memory_image_slot( - &self, - instance_index: usize, - memory_index: DefinedMemoryIndex, - slot: MemoryImageSlot, - ) { - assert!(!slot.is_dirty()); - let idx = instance_index * self.max_memories + (memory_index.as_u32() as usize); - *self.image_slots[idx].lock().unwrap() = Some(slot); - } - - /// Resets all the images for the instance index slot specified to clear out - /// any prior mappings. - /// - /// This is used when a `Module` is dropped at the `wasmtime` layer to clear - /// out any remaining mappings and ensure that its memfd backing, if any, is - /// removed from the address space to avoid lingering references to it. - fn clear_images(&self, instance_index: usize) { - for i in 0..self.max_memories { - let index = DefinedMemoryIndex::from_u32(i as u32); - - // Clear the image from the slot and, if successful, return it back - // to our state. Note that on failure here the whole slot will get - // paved over with an anonymous mapping. - let mut slot = self.take_memory_image_slot(instance_index, index); - if slot.remove_image().is_ok() { - self.return_memory_image_slot(instance_index, index, slot); - } - } - } -} - -impl Drop for MemoryPool { - fn drop(&mut self) { - // Clear the `clear_no_drop` flag (i.e., ask to *not* clear on - // drop) for all slots, and then drop them here. This is - // valid because the one `Mmap` that covers the whole region - // can just do its one munmap. - for mut slot in std::mem::take(&mut self.image_slots) { - if let Some(slot) = slot.get_mut().unwrap() { - slot.no_clear_on_drop(); - } - } - } -} - -/// Represents a pool of WebAssembly tables. -/// -/// Each instance index into the pool returns an iterator over the base addresses -/// of the instance's tables. -#[derive(Debug)] -struct TablePool { - mapping: Mmap, - table_size: usize, - max_tables: usize, - max_instances: usize, - page_size: usize, - max_elements: u32, -} - -impl TablePool { - fn new(instance_limits: &InstanceLimits) -> Result { - let page_size = crate::page_size(); - - let table_size = round_up_to_pow2( - mem::size_of::<*mut u8>() - .checked_mul(instance_limits.table_elements as usize) - .ok_or_else(|| anyhow!("table size exceeds addressable memory"))?, - page_size, - ); - - let max_instances = instance_limits.count as usize; - let max_tables = instance_limits.tables as usize; - - let allocation_size = table_size - .checked_mul(max_tables) - .and_then(|c| c.checked_mul(max_instances)) - .ok_or_else(|| anyhow!("total size of instance tables exceeds addressable memory"))?; - - let mapping = Mmap::accessible_reserved(allocation_size, allocation_size) - .context("failed to create table pool mapping")?; - - Ok(Self { - mapping, - table_size, - max_tables, - max_instances, - page_size, - max_elements: instance_limits.table_elements, - }) - } - - fn get(&self, instance_index: usize) -> impl Iterator { - assert!(instance_index < self.max_instances); - - let base: *mut u8 = unsafe { - self.mapping - .as_ptr() - .add(instance_index * self.table_size * self.max_tables) - .cast_mut() - }; - - let size = self.table_size; - (0..self.max_tables).map(move |i| unsafe { base.add(i * size) }) - } -} - -/// Represents a pool of execution stacks (used for the async fiber implementation). -/// -/// Each index into the pool represents a single execution stack. The maximum number of -/// stacks is the same as the maximum number of instances. -/// -/// As stacks grow downwards, each stack starts (lowest address) with a guard page -/// that can be used to detect stack overflow. -/// -/// The top of the stack (starting stack pointer) is returned when a stack is allocated -/// from the pool. -#[cfg(all(feature = "async", unix, not(miri)))] -#[derive(Debug)] -struct StackPool { - mapping: Mmap, - stack_size: usize, - max_instances: usize, - page_size: usize, - index_allocator: ModuleAffinityIndexAllocator, - async_stack_zeroing: bool, - async_stack_keep_resident: usize, -} - -#[cfg(all(feature = "async", unix, not(miri)))] -impl StackPool { - fn new(config: &PoolingInstanceAllocatorConfig) -> Result { - use rustix::mm::{mprotect, MprotectFlags}; - - let page_size = crate::page_size(); - - // Add a page to the stack size for the guard page when using fiber stacks - let stack_size = if config.stack_size == 0 { - 0 - } else { - round_up_to_pow2(config.stack_size, page_size) - .checked_add(page_size) - .ok_or_else(|| anyhow!("stack size exceeds addressable memory"))? - }; - - let max_instances = config.limits.count as usize; - - let allocation_size = stack_size - .checked_mul(max_instances) - .ok_or_else(|| anyhow!("total size of execution stacks exceeds addressable memory"))?; - - let mapping = Mmap::accessible_reserved(allocation_size, allocation_size) - .context("failed to create stack pool mapping")?; - - // Set up the stack guard pages - if allocation_size > 0 { - unsafe { - for i in 0..max_instances { - // Make the stack guard page inaccessible - let bottom_of_stack = mapping.as_ptr().add(i * stack_size).cast_mut(); - mprotect(bottom_of_stack.cast(), page_size, MprotectFlags::empty()) - .context("failed to protect stack guard page")?; - } - } - } - - Ok(Self { - mapping, - stack_size, - max_instances, - page_size, - async_stack_zeroing: config.async_stack_zeroing, - async_stack_keep_resident: config.async_stack_keep_resident, - // Note that `max_unused_warm_slots` is set to zero since stacks - // have no affinity so there's no need to keep intentionally unused - // warm slots around. - index_allocator: ModuleAffinityIndexAllocator::new(config.limits.count, 0), - }) - } - - fn allocate(&self) -> Result { - if self.stack_size == 0 { - bail!("pooling allocator not configured to enable fiber stack allocation"); - } - - let index = self - .index_allocator - .alloc(None) - .ok_or_else(|| { - anyhow!( - "maximum concurrent fiber limit of {} reached", - self.max_instances - ) - })? - .index(); - - assert!(index < self.max_instances); - - unsafe { - // Remove the guard page from the size - let size_without_guard = self.stack_size - self.page_size; - - let bottom_of_stack = self - .mapping - .as_ptr() - .add((index * self.stack_size) + self.page_size) - .cast_mut(); - - commit_stack_pages(bottom_of_stack, size_without_guard)?; - - let stack = - wasmtime_fiber::FiberStack::from_raw_parts(bottom_of_stack, size_without_guard)?; - Ok(stack) - } - } - - fn deallocate(&self, stack: &wasmtime_fiber::FiberStack) { - let top = stack - .top() - .expect("fiber stack not allocated from the pool") as usize; - - let base = self.mapping.as_ptr() as usize; - let len = self.mapping.len(); - assert!( - top > base && top <= (base + len), - "fiber stack top pointer not in range" - ); - - // Remove the guard page from the size - let stack_size = self.stack_size - self.page_size; - let bottom_of_stack = top - stack_size; - let start_of_stack = bottom_of_stack - self.page_size; - assert!(start_of_stack >= base && start_of_stack < (base + len)); - assert!((start_of_stack - base) % self.stack_size == 0); - - let index = (start_of_stack - base) / self.stack_size; - assert!(index < self.max_instances); - - if self.async_stack_zeroing { - self.zero_stack(bottom_of_stack, stack_size); - } - - self.index_allocator.free(SlotId(index as u32)); - } - - fn zero_stack(&self, bottom: usize, size: usize) { - // Manually zero the top of the stack to keep the pages resident in - // memory and avoid future page faults. Use the system to deallocate - // pages past this. This hopefully strikes a reasonable balance between: - // - // * memset for the whole range is probably expensive - // * madvise for the whole range incurs expensive future page faults - // * most threads probably don't use most of the stack anyway - let size_to_memset = size.min(self.async_stack_keep_resident); - unsafe { - std::ptr::write_bytes( - (bottom + size - size_to_memset) as *mut u8, - 0, - size_to_memset, - ); - } - - // Use the system to reset remaining stack pages to zero. - reset_stack_pages_to_zero(bottom as _, size - size_to_memset).unwrap(); - } -} - /// Configuration options for the pooling instance allocator supplied at /// construction. #[derive(Copy, Clone, Debug)] @@ -570,13 +185,20 @@ impl Default for PoolingInstanceAllocatorConfig { /// Note: the resource pools are manually dropped so that the fault handler terminates correctly. #[derive(Debug)] pub struct PoolingInstanceAllocator { - instance_size: usize, - max_instances: usize, - index_allocator: ModuleAffinityIndexAllocator, + limits: InstanceLimits, + + // The number of live core module and component instances at any given + // time. Note that this can temporarily go over the configured limit. This + // doesn't mean we have actually overshot, but that we attempted to allocate + // a new instance and incremented the counter, we've seen (or are about to + // see) that the counter is beyond the configured threshold, and are going + // to decrement the counter and return an error but haven't done so yet. See + // the increment trait methods for more details. + live_core_instances: AtomicU64, + live_component_instances: AtomicU64, + memories: MemoryPool, tables: TablePool, - linear_memory_keep_resident: usize, - table_keep_resident: usize, #[cfg(all(feature = "async", unix, not(miri)))] stacks: StackPool, @@ -587,20 +209,12 @@ pub struct PoolingInstanceAllocator { impl PoolingInstanceAllocator { /// Creates a new pooling instance allocator with the given strategy and limits. pub fn new(config: &PoolingInstanceAllocatorConfig, tunables: &Tunables) -> Result { - if config.limits.count == 0 { - bail!("the instance count limit cannot be zero"); - } - - let max_instances = config.limits.count as usize; - Ok(Self { - instance_size: round_up_to_pow2(config.limits.size, mem::align_of::()), - max_instances, - index_allocator: ModuleAffinityIndexAllocator::new(config.limits.count, config.max_unused_warm_slots), - memories: MemoryPool::new(&config.limits, tunables)?, - tables: TablePool::new(&config.limits)?, - linear_memory_keep_resident: config.linear_memory_keep_resident, - table_keep_resident: config.table_keep_resident, + limits: config.limits, + live_component_instances: AtomicU64::new(0), + live_core_instances: AtomicU64::new(0), + memories: MemoryPool::new(config, tunables)?, + tables: TablePool::new(config)?, #[cfg(all(feature = "async", unix, not(miri)))] stacks: StackPool::new(config)?, #[cfg(all(feature = "async", windows))] @@ -608,81 +222,21 @@ impl PoolingInstanceAllocator { }) } - fn reset_table_pages_to_zero(&self, base: *mut u8, size: usize) -> Result<()> { - let size_to_memset = size.min(self.table_keep_resident); - unsafe { - std::ptr::write_bytes(base, 0, size_to_memset); - decommit_table_pages(base.add(size_to_memset), size - size_to_memset) - .context("failed to decommit table page")?; - } - Ok(()) + fn core_instance_size(&self) -> usize { + round_up_to_pow2(self.limits.core_instance_size, mem::align_of::()) } fn validate_table_plans(&self, module: &Module) -> Result<()> { - let tables = module.table_plans.len() - module.num_imported_tables; - if tables > self.tables.max_tables { - bail!( - "defined tables count of {} exceeds the limit of {}", - tables, - self.tables.max_tables, - ); - } - - for (i, plan) in module.table_plans.iter().skip(module.num_imported_tables) { - if plan.table.minimum > self.tables.max_elements { - bail!( - "table index {} has a minimum element size of {} which exceeds the limit of {}", - i.as_u32(), - plan.table.minimum, - self.tables.max_elements, - ); - } - } - Ok(()) + self.tables.validate(module) } fn validate_memory_plans(&self, module: &Module) -> Result<()> { - let memories = module.memory_plans.len() - module.num_imported_memories; - if memories > self.memories.max_memories { - bail!( - "defined memories count of {} exceeds the limit of {}", - memories, - self.memories.max_memories, - ); - } - - for (i, plan) in module - .memory_plans - .iter() - .skip(module.num_imported_memories) - { - match plan.style { - MemoryStyle::Static { bound } => { - if (self.memories.memory_size as u64) < bound { - bail!( - "memory size allocated per-memory is too small to \ - satisfy static bound of {bound:#x} pages" - ); - } - } - MemoryStyle::Dynamic { .. } => {} - } - let max = self.memories.max_accessible / (WASM_PAGE_SIZE as usize); - if plan.memory.minimum > (max as u64) { - bail!( - "memory index {} has a minimum page size of {} which exceeds the limit of {}", - i.as_u32(), - plan.memory.minimum, - max, - ); - } - } - Ok(()) + self.memories.validate(module) } - fn validate_instance_size(&self, offsets: &VMOffsets) -> Result<()> { + fn validate_core_instance_size(&self, offsets: &VMOffsets) -> Result<()> { let layout = Instance::alloc_layout(offsets); - if layout.size() <= self.instance_size { + if layout.size() <= self.core_instance_size() { return Ok(()); } @@ -698,7 +252,7 @@ impl PoolingInstanceAllocator { requires {} bytes which exceeds the configured maximum \ of {} bytes; breakdown of allocation requirement:\n\n", layout.size(), - self.instance_size, + self.core_instance_size(), ); let mut remaining = layout.size(); @@ -736,161 +290,168 @@ impl PoolingInstanceAllocator { bail!("{}", message) } -} - -unsafe impl InstanceAllocator for PoolingInstanceAllocator { - fn validate(&self, module: &Module, offsets: &VMOffsets) -> Result<()> { - self.validate_memory_plans(module)?; - self.validate_table_plans(module)?; - self.validate_instance_size(offsets)?; - - Ok(()) - } - fn allocate_index(&self, req: &InstanceAllocationRequest) -> Result { - self.index_allocator - .alloc(req.runtime_info.unique_id()) - .map(|id| id.index()) - .ok_or_else(|| { - anyhow!( - "maximum concurrent instance limit of {} reached", - self.max_instances - ) - }) - } + #[cfg(feature = "component-model")] + fn validate_component_instance_size( + &self, + offsets: &VMComponentOffsets, + ) -> Result<()> { + if usize::try_from(offsets.size_of_vmctx()).unwrap() <= self.limits.component_instance_size + { + return Ok(()); + } - fn deallocate_index(&self, index: usize) { - self.index_allocator.free(SlotId(index as u32)); + // TODO: Add context with detailed accounting of what makes up all the + // `VMComponentContext`'s space like we do for module instances. + bail!( + "instance allocation for this component requires {} bytes of `VMComponentContext` \ + space which exceeds the configured maximum of {} bytes", + offsets.size_of_vmctx(), + self.limits.component_instance_size + ) } +} - fn allocate_memories( +unsafe impl InstanceAllocatorImpl for PoolingInstanceAllocator { + #[cfg(feature = "component-model")] + fn validate_component_impl<'a>( &self, - index: usize, - req: &mut InstanceAllocationRequest, - memories: &mut PrimaryMap, + component: &Component, + offsets: &VMComponentOffsets, + get_module: &'a dyn Fn(StaticModuleIndex) -> &'a Module, ) -> Result<()> { - let module = req.runtime_info.module(); - - self.validate_memory_plans(module)?; - - for (memory_index, plan) in module - .memory_plans - .iter() - .skip(module.num_imported_memories) - { - let defined_index = module - .defined_memory_index(memory_index) - .expect("should be a defined memory since we skipped imported ones"); - - // Double-check that the runtime requirements of the memory are - // satisfied by the configuration of this pooling allocator. This - // should be returned as an error through `validate_memory_plans` - // but double-check here to be sure. - match plan.style { - MemoryStyle::Static { bound } => { - let bound = bound * u64::from(WASM_PAGE_SIZE); - assert!(bound <= (self.memories.memory_size as u64)); + self.validate_component_instance_size(offsets)?; + + let mut num_core_instances = 0; + let mut num_memories = 0; + let mut num_tables = 0; + for init in &component.initializers { + use wasmtime_environ::component::GlobalInitializer::*; + use wasmtime_environ::component::InstantiateModule; + match init { + InstantiateModule(InstantiateModule::Import(_, _)) => { + num_core_instances += 1; + // Can't statically account for the total vmctx size, number + // of memories, and number of tables in this component. + } + InstantiateModule(InstantiateModule::Static(static_module_index, _)) => { + let module = get_module(*static_module_index); + let offsets = VMOffsets::new(HostPtr, &module); + self.validate_module_impl(module, &offsets)?; + num_core_instances += 1; + num_memories += module.memory_plans.len() - module.num_imported_memories; + num_tables += module.table_plans.len() - module.num_imported_tables; } - MemoryStyle::Dynamic { .. } => {} + LowerImport { .. } + | ExtractMemory(_) + | ExtractRealloc(_) + | ExtractPostReturn(_) + | Resource(_) => {} } + } - let base_ptr = self.memories.get_base(index, defined_index); - let base_capacity = self.memories.max_accessible; - - let mut slot = self.memories.take_memory_image_slot(index, defined_index); - let image = req.runtime_info.memory_image(defined_index)?; - let initial_size = plan.memory.minimum * WASM_PAGE_SIZE as u64; - - // If instantiation fails, we can propagate the error - // upward and drop the slot. This will cause the Drop - // handler to attempt to map the range with PROT_NONE - // memory, to reserve the space while releasing any - // stale mappings. The next use of this slot will then - // create a new slot that will try to map over - // this, returning errors as well if the mapping - // errors persist. The unmap-on-drop is best effort; - // if it fails, then we can still soundly continue - // using the rest of the pool and allowing the rest of - // the process to continue, because we never perform a - // mmap that would leave an open space for someone - // else to come in and map something. - slot.instantiate(initial_size as usize, image, &plan)?; - - memories.push(Memory::new_static( - plan, - base_ptr, - base_capacity, - slot, - self.memories.memory_and_guard_size, - unsafe { &mut *req.store.get().unwrap() }, - )?); + if num_core_instances + > usize::try_from(self.limits.max_core_instances_per_component).unwrap() + { + bail!( + "The component transitively contains {num_core_instances} core module instances, \ + which exceeds the configured maximum of {}", + self.limits.max_core_instances_per_component + ); } - Ok(()) - } + if num_memories > usize::try_from(self.limits.max_memories_per_component).unwrap() { + bail!( + "The component transitively contains {num_memories} Wasm linear memories, which \ + exceeds the configured maximum of {}", + self.limits.max_memories_per_component + ); + } - fn deallocate_memories(&self, index: usize, mems: &mut PrimaryMap) { - // Decommit any linear memories that were used. - for (def_mem_idx, memory) in mem::take(mems) { - let mut image = memory.unwrap_static_image(); - // 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 - // reservation. - if image - .clear_and_remain_ready(self.linear_memory_keep_resident) - .is_ok() - { - self.memories - .return_memory_image_slot(index, def_mem_idx, image); - } + if num_tables > usize::try_from(self.limits.max_tables_per_component).unwrap() { + bail!( + "The component transitively contains {num_tables} tables, which exceeds the \ + configured maximum of {}", + self.limits.max_tables_per_component + ); } - } - fn allocate_tables( - &self, - index: usize, - req: &mut InstanceAllocationRequest, - tables: &mut PrimaryMap, - ) -> Result<()> { - let module = req.runtime_info.module(); + Ok(()) + } + fn validate_module_impl(&self, module: &Module, offsets: &VMOffsets) -> Result<()> { + self.validate_memory_plans(module)?; self.validate_table_plans(module)?; + self.validate_core_instance_size(offsets)?; + Ok(()) + } - let mut bases = self.tables.get(index); - for (_, plan) in module.table_plans.iter().skip(module.num_imported_tables) { - let base = bases.next().unwrap() as _; + fn increment_component_instance_count(&self) -> Result<()> { + let old_count = self.live_component_instances.fetch_add(1, Ordering::AcqRel); + if old_count >= u64::from(self.limits.total_component_instances) { + self.decrement_component_instance_count(); + bail!( + "maximum concurrent component instance limit of {} reached", + self.limits.total_component_instances + ); + } + Ok(()) + } - commit_table_pages( - base as *mut u8, - self.tables.max_elements as usize * mem::size_of::<*mut u8>(), - )?; + fn decrement_component_instance_count(&self) { + self.live_component_instances.fetch_sub(1, Ordering::AcqRel); + } - tables.push(Table::new_static( - plan, - unsafe { std::slice::from_raw_parts_mut(base, self.tables.max_elements as usize) }, - unsafe { &mut *req.store.get().unwrap() }, - )?); + fn increment_core_instance_count(&self) -> Result<()> { + let old_count = self.live_core_instances.fetch_add(1, Ordering::AcqRel); + if old_count >= u64::from(self.limits.total_core_instances) { + self.decrement_core_instance_count(); + bail!( + "maximum concurrent core instance limit of {} reached", + self.limits.total_core_instances + ); } - Ok(()) } - fn deallocate_tables(&self, index: usize, tables: &mut PrimaryMap) { - // Decommit any tables that were used - for (table, base) in tables.values_mut().zip(self.tables.get(index)) { - let table = mem::take(table); - assert!(table.is_static()); + fn decrement_core_instance_count(&self) { + self.live_core_instances.fetch_sub(1, Ordering::AcqRel); + } - let size = round_up_to_pow2( - table.size() as usize * mem::size_of::<*mut u8>(), - self.tables.page_size, - ); + unsafe fn allocate_memory( + &self, + request: &mut InstanceAllocationRequest, + memory_plan: &MemoryPlan, + memory_index: DefinedMemoryIndex, + ) -> Result<(MemoryAllocationIndex, Memory)> { + self.memories.allocate(request, memory_plan, memory_index) + } - drop(table); - self.reset_table_pages_to_zero(base, size) - .expect("failed to decommit table pages"); - } + unsafe fn deallocate_memory( + &self, + _memory_index: DefinedMemoryIndex, + allocation_index: MemoryAllocationIndex, + memory: Memory, + ) { + self.memories.deallocate(allocation_index, memory); + } + + unsafe fn allocate_table( + &self, + request: &mut InstanceAllocationRequest, + table_plan: &TablePlan, + _table_index: DefinedTableIndex, + ) -> Result<(super::TableAllocationIndex, Table)> { + self.tables.allocate(request, table_plan) + } + + unsafe fn deallocate_table( + &self, + _table_index: DefinedTableIndex, + allocation_index: TableAllocationIndex, + table: Table, + ) { + self.tables.deallocate(allocation_index, table); } #[cfg(feature = "async")] @@ -933,328 +494,19 @@ unsafe impl InstanceAllocator for PoolingInstanceAllocator { } fn purge_module(&self, module: CompiledModuleId) { - // Purging everything related to `module` primarily means clearing out - // all of its memory images present in the virtual address space. Go - // through the index allocator for slots affine to `module` and reset - // them, freeing up the index when we're done. - // - // Note that this is only called when the specified `module` won't be - // allocated further (the module is being dropped) so this shouldn't hit - // any sort of infinite loop since this should be the final operation - // working with `module`. - while let Some(index) = self.index_allocator.alloc_affine_and_clear_affinity(module) { - self.memories.clear_images(index.index()); - self.index_allocator.free(index); - } + self.memories.purge_module(module); } } #[cfg(test)] mod test { use super::*; - use crate::{ - CompiledModuleId, Imports, MemoryImage, ModuleRuntimeInfo, StorePtr, VMSharedSignatureIndex, - }; - use std::{ptr::NonNull, sync::Arc}; - use wasmtime_environ::{DefinedFuncIndex, DefinedMemoryIndex}; - - pub(crate) fn empty_runtime_info( - module: Arc, - ) -> Arc { - struct RuntimeInfo(Arc, VMOffsets); - - impl ModuleRuntimeInfo for RuntimeInfo { - fn module(&self) -> &Arc { - &self.0 - } - fn function(&self, _: DefinedFuncIndex) -> NonNull { - unimplemented!() - } - fn array_to_wasm_trampoline( - &self, - _: DefinedFuncIndex, - ) -> Option { - unimplemented!() - } - fn native_to_wasm_trampoline( - &self, - _: DefinedFuncIndex, - ) -> Option> { - unimplemented!() - } - fn wasm_to_native_trampoline( - &self, - _: VMSharedSignatureIndex, - ) -> Option> { - unimplemented!() - } - fn memory_image( - &self, - _: DefinedMemoryIndex, - ) -> anyhow::Result>> { - Ok(None) - } - - fn unique_id(&self) -> Option { - None - } - fn wasm_data(&self) -> &[u8] { - &[] - } - fn signature_ids(&self) -> &[VMSharedSignatureIndex] { - &[] - } - fn offsets(&self) -> &VMOffsets { - &self.1 - } - } - - let offsets = VMOffsets::new(HostPtr, &module); - Arc::new(RuntimeInfo(module, offsets)) - } - - #[cfg(target_pointer_width = "64")] - #[test] - fn test_instance_pool() -> Result<()> { - let mut config = PoolingInstanceAllocatorConfig::default(); - config.max_unused_warm_slots = 0; - config.limits = InstanceLimits { - count: 3, - tables: 1, - memories: 1, - table_elements: 10, - size: 1000, - memory_pages: 1, - ..Default::default() - }; - - let instances = PoolingInstanceAllocator::new( - &config, - &Tunables { - static_memory_bound: 1, - ..Tunables::default() - }, - )?; - - assert_eq!(instances.instance_size, 1008); // round 1000 up to alignment - assert_eq!(instances.max_instances, 3); - - assert_eq!(instances.index_allocator.testing_freelist(), []); - - let mut handles = Vec::new(); - let module = Arc::new(Module::default()); - - for _ in (0..3).rev() { - handles.push( - instances - .allocate(InstanceAllocationRequest { - runtime_info: &empty_runtime_info(module.clone()), - imports: Imports { - functions: &[], - tables: &[], - memories: &[], - globals: &[], - }, - host_state: Box::new(()), - store: StorePtr::empty(), - wmemcheck: false, - }) - .expect("allocation should succeed"), - ); - } - - assert_eq!(instances.index_allocator.testing_freelist(), []); - - match instances.allocate(InstanceAllocationRequest { - runtime_info: &empty_runtime_info(module), - imports: Imports { - functions: &[], - tables: &[], - memories: &[], - globals: &[], - }, - host_state: Box::new(()), - store: StorePtr::empty(), - wmemcheck: false, - }) { - Err(_) => {} - _ => panic!("unexpected error"), - }; - - for mut handle in handles.drain(..) { - instances.deallocate(&mut handle); - } - - assert_eq!( - instances.index_allocator.testing_freelist(), - [SlotId(0), SlotId(1), SlotId(2)] - ); - - Ok(()) - } - - #[cfg(target_pointer_width = "64")] - #[test] - fn test_memory_pool() -> Result<()> { - let pool = MemoryPool::new( - &InstanceLimits { - count: 5, - tables: 0, - memories: 3, - table_elements: 0, - memory_pages: 1, - ..Default::default() - }, - &Tunables { - static_memory_bound: 1, - static_memory_offset_guard_size: 0, - ..Tunables::default() - }, - )?; - - assert_eq!(pool.memory_and_guard_size, WASM_PAGE_SIZE as usize); - assert_eq!(pool.max_memories, 3); - assert_eq!(pool.max_instances, 5); - assert_eq!(pool.max_accessible, WASM_PAGE_SIZE as usize); - - let base = pool.mapping.as_ptr() as usize; - - for i in 0..5 { - let mut iter = pool.get(i); - - for j in 0..3 { - assert_eq!( - iter.next().unwrap() as usize - base, - ((i * 3) + j) * pool.memory_and_guard_size - ); - } - - assert_eq!(iter.next(), None); - } - - Ok(()) - } - - #[cfg(target_pointer_width = "64")] - #[test] - fn test_table_pool() -> Result<()> { - let pool = TablePool::new(&InstanceLimits { - count: 7, - table_elements: 100, - memory_pages: 0, - tables: 4, - memories: 0, - ..Default::default() - })?; - - let host_page_size = crate::page_size(); - - assert_eq!(pool.table_size, host_page_size); - assert_eq!(pool.max_tables, 4); - assert_eq!(pool.max_instances, 7); - assert_eq!(pool.page_size, host_page_size); - assert_eq!(pool.max_elements, 100); - - let base = pool.mapping.as_ptr() as usize; - - for i in 0..7 { - let mut iter = pool.get(i); - - for j in 0..4 { - assert_eq!( - iter.next().unwrap() as usize - base, - ((i * 4) + j) * pool.table_size - ); - } - - assert_eq!(iter.next(), None); - } - - Ok(()) - } - - #[cfg(all(unix, target_pointer_width = "64", feature = "async", not(miri)))] - #[test] - fn test_stack_pool() -> Result<()> { - let config = PoolingInstanceAllocatorConfig { - limits: InstanceLimits { - count: 10, - ..Default::default() - }, - stack_size: 1, - async_stack_zeroing: true, - ..PoolingInstanceAllocatorConfig::default() - }; - let pool = StackPool::new(&config)?; - - let native_page_size = crate::page_size(); - assert_eq!(pool.stack_size, 2 * native_page_size); - assert_eq!(pool.max_instances, 10); - assert_eq!(pool.page_size, native_page_size); - - assert_eq!(pool.index_allocator.testing_freelist(), []); - - let base = pool.mapping.as_ptr() as usize; - - let mut stacks = Vec::new(); - for i in 0..10 { - let stack = pool.allocate().expect("allocation should succeed"); - assert_eq!( - ((stack.top().unwrap() as usize - base) / pool.stack_size) - 1, - i - ); - stacks.push(stack); - } - - assert_eq!(pool.index_allocator.testing_freelist(), []); - - pool.allocate().unwrap_err(); - - for stack in stacks { - pool.deallocate(&stack); - } - - assert_eq!( - pool.index_allocator.testing_freelist(), - [ - SlotId(0), - SlotId(1), - SlotId(2), - SlotId(3), - SlotId(4), - SlotId(5), - SlotId(6), - SlotId(7), - SlotId(8), - SlotId(9) - ], - ); - - Ok(()) - } - - #[test] - fn test_pooling_allocator_with_zero_instance_count() { - let config = PoolingInstanceAllocatorConfig { - limits: InstanceLimits { - count: 0, - ..Default::default() - }, - ..PoolingInstanceAllocatorConfig::default() - }; - assert_eq!( - PoolingInstanceAllocator::new(&config, &Tunables::default(),) - .map_err(|e| e.to_string()) - .expect_err("expected a failure constructing instance allocator"), - "the instance count limit cannot be zero" - ); - } #[test] fn test_pooling_allocator_with_memory_pages_exceeded() { let config = PoolingInstanceAllocatorConfig { limits: InstanceLimits { - count: 1, + total_memories: 1, memory_pages: 0x10001, ..Default::default() }, @@ -1274,39 +526,15 @@ mod test { ); } - #[test] - fn test_pooling_allocator_with_reservation_size_exceeded() { - let config = PoolingInstanceAllocatorConfig { - limits: InstanceLimits { - count: 1, - memory_pages: 2, - ..Default::default() - }, - ..PoolingInstanceAllocatorConfig::default() - }; - let pool = PoolingInstanceAllocator::new( - &config, - &Tunables { - static_memory_bound: 1, - static_memory_offset_guard_size: 0, - ..Tunables::default() - }, - ) - .unwrap(); - assert_eq!(pool.memories.memory_size, 2 * 65536); - } - #[cfg(all(unix, target_pointer_width = "64", feature = "async", not(miri)))] #[test] fn test_stack_zeroed() -> Result<()> { let config = PoolingInstanceAllocatorConfig { max_unused_warm_slots: 0, limits: InstanceLimits { - count: 1, - table_elements: 0, - memory_pages: 0, - tables: 0, - memories: 0, + total_stacks: 1, + total_memories: 0, + total_tables: 0, ..Default::default() }, stack_size: 128, @@ -1338,11 +566,9 @@ mod test { let config = PoolingInstanceAllocatorConfig { max_unused_warm_slots: 0, limits: InstanceLimits { - count: 1, - table_elements: 0, - memory_pages: 0, - tables: 0, - memories: 0, + total_stacks: 1, + total_memories: 0, + total_tables: 0, ..Default::default() }, stack_size: 128, diff --git a/crates/runtime/src/instance/allocator/pooling/index_allocator.rs b/crates/runtime/src/instance/allocator/pooling/index_allocator.rs index 43433185a3dd..ec4dbf1060f8 100644 --- a/crates/runtime/src/instance/allocator/pooling/index_allocator.rs +++ b/crates/runtime/src/instance/allocator/pooling/index_allocator.rs @@ -4,6 +4,7 @@ use crate::CompiledModuleId; use std::collections::hash_map::{Entry, HashMap}; use std::mem; use std::sync::Mutex; +use wasmtime_environ::DefinedMemoryIndex; /// A slot index. #[derive(Hash, Clone, Copy, Debug, PartialEq, Eq)] @@ -36,6 +37,11 @@ impl SimpleIndexAllocator { pub(crate) fn free(&self, index: SlotId) { self.0.free(index); } + + #[cfg(test)] + pub(crate) fn testing_freelist(&self) -> Vec { + self.0.testing_freelist() + } } /// An index allocator that has configurable affinity between slots and modules @@ -45,7 +51,7 @@ pub struct ModuleAffinityIndexAllocator(Mutex); #[derive(Debug)] struct Inner { - /// Maximum number of "unused warm slots" which will be allowed during + /// Maximum number of "unused warm slots" which will be allowed during /// allocation. /// /// This is a user-configurable knob which can be used to influence the @@ -80,7 +86,7 @@ struct Inner { /// /// The `List` here is appended to during deallocation and removal happens /// from the tail during allocation. - module_affine: HashMap, + module_affine: HashMap<(CompiledModuleId, DefinedMemoryIndex), List>, } /// A helper "linked list" data structure which is based on indices. @@ -100,8 +106,8 @@ struct Link { #[derive(Clone, Debug)] enum SlotState { - /// This slot is currently in use and is affine to the specified module. - Used(Option), + /// This slot is currently in use and is affine to the specified module's memory. + Used(Option<(CompiledModuleId, DefinedMemoryIndex)>), /// This slot is not currently used, and has never been used. UnusedCold, @@ -125,7 +131,7 @@ impl SlotState { #[derive(Default, Copy, Clone, Debug)] struct Unused { /// Which module this slot was historically affine to, if any. - affinity: Option, + affinity: Option<(CompiledModuleId, DefinedMemoryIndex)>, /// Metadata about the linked list for all slots affine to `affinity`. affine_list_link: Link, @@ -156,8 +162,11 @@ impl ModuleAffinityIndexAllocator { /// affinity request if the allocation strategy supports it. /// /// Returns `None` if no more slots are available. - pub fn alloc(&self, module_id: Option) -> Option { - self._alloc(module_id, AllocMode::AnySlot) + pub fn alloc( + &self, + for_memory: Option<(CompiledModuleId, DefinedMemoryIndex)>, + ) -> Option { + self._alloc(for_memory, AllocMode::AnySlot) } /// Attempts to allocate a guaranteed-affine slot to the module `id` @@ -167,18 +176,29 @@ impl ModuleAffinityIndexAllocator { /// this slot will not record the affinity to `id`, instead simply listing /// it as taken. This is intended to be used for clearing out all affine /// slots to a module. - pub fn alloc_affine_and_clear_affinity(&self, module_id: CompiledModuleId) -> Option { - self._alloc(Some(module_id), AllocMode::ForceAffineAndClear) + pub fn alloc_affine_and_clear_affinity( + &self, + module_id: CompiledModuleId, + memory_index: DefinedMemoryIndex, + ) -> Option { + self._alloc( + Some((module_id, memory_index)), + AllocMode::ForceAffineAndClear, + ) } - fn _alloc(&self, module_id: Option, mode: AllocMode) -> Option { + fn _alloc( + &self, + for_memory: Option<(CompiledModuleId, DefinedMemoryIndex)>, + mode: AllocMode, + ) -> Option { let mut inner = self.0.lock().unwrap(); let inner = &mut *inner; // As a first-pass always attempt an affine allocation. This will // succeed if any slots are considered affine to `module_id` (if it's // specified). Failing that something else is attempted to be chosen. - let slot_id = inner.pick_affine(module_id).or_else(|| { + let slot_id = inner.pick_affine(for_memory).or_else(|| { match mode { // If any slot is requested then this is a normal instantiation // looking for an index. Without any affine candidates there are @@ -221,7 +241,7 @@ impl ModuleAffinityIndexAllocator { inner.slot_state[slot_id.index()] = SlotState::Used(match mode { AllocMode::ForceAffineAndClear => None, - AllocMode::AnySlot => module_id, + AllocMode::AnySlot => for_memory, }); Some(slot_id) @@ -230,8 +250,8 @@ impl ModuleAffinityIndexAllocator { pub(crate) fn free(&self, index: SlotId) { let mut inner = self.0.lock().unwrap(); let inner = &mut *inner; - let module = match inner.slot_state[index.index()] { - SlotState::Used(module) => module, + let module_memory = match inner.slot_state[index.index()] { + SlotState::Used(module_memory) => module_memory, _ => unreachable!(), }; @@ -243,7 +263,7 @@ impl ModuleAffinityIndexAllocator { .warm .append(index, &mut inner.slot_state, |s| &mut s.unused_list_link); - let affine_list_link = match module { + let affine_list_link = match module_memory { // If this slot is affine to a particular module then append this // index to the linked list for the affine module. Otherwise insert // a new one-element linked list. @@ -262,7 +282,7 @@ impl ModuleAffinityIndexAllocator { }; inner.slot_state[index.index()] = SlotState::UnusedWarm(Unused { - affinity: module, + affinity: module_memory, affine_list_link, unused_list_link, }); @@ -282,7 +302,9 @@ impl ModuleAffinityIndexAllocator { /// For testing only, get the list of all modules with at least /// one slot with affinity for that module. #[cfg(test)] - pub(crate) fn testing_module_affinity_list(&self) -> Vec { + pub(crate) fn testing_module_affinity_list( + &self, + ) -> Vec<(CompiledModuleId, DefinedMemoryIndex)> { let inner = self.0.lock().unwrap(); inner.module_affine.keys().copied().collect() } @@ -291,11 +313,14 @@ impl ModuleAffinityIndexAllocator { impl Inner { /// Attempts to allocate a slot already affine to `id`, returning `None` if /// `id` is `None` or if there are no affine slots. - fn pick_affine(&mut self, module_id: Option) -> Option { + fn pick_affine( + &mut self, + for_memory: Option<(CompiledModuleId, DefinedMemoryIndex)>, + ) -> Option { // Note that the `tail` is chosen here of the affine list as it's the // most recently used, which for affine allocations is what we want -- // maximizing temporal reuse. - let ret = self.module_affine.get(&module_id?)?.tail?; + let ret = self.module_affine.get(&for_memory?)?.tail?; self.remove(ret); Some(ret) } @@ -434,8 +459,9 @@ impl List { #[cfg(test)] mod test { - use super::{ModuleAffinityIndexAllocator, SlotId}; + use super::*; use crate::CompiledModuleIdAllocator; + use wasmtime_environ::EntityRef; #[test] fn test_next_available_allocation_strategy() { @@ -451,8 +477,8 @@ mod test { #[test] fn test_affinity_allocation_strategy() { let id_alloc = CompiledModuleIdAllocator::new(); - let id1 = id_alloc.alloc(); - let id2 = id_alloc.alloc(); + let id1 = (id_alloc.alloc(), DefinedMemoryIndex::new(0)); + let id2 = (id_alloc.alloc(), DefinedMemoryIndex::new(0)); let state = ModuleAffinityIndexAllocator::new(100, 100); let index1 = state.alloc(Some(id1)).unwrap(); @@ -505,17 +531,25 @@ mod test { fn clear_affine() { let id_alloc = CompiledModuleIdAllocator::new(); let id = id_alloc.alloc(); + let memory_index = DefinedMemoryIndex::new(0); for max_unused_warm_slots in [0, 1, 2] { let state = ModuleAffinityIndexAllocator::new(100, max_unused_warm_slots); - let index1 = state.alloc(Some(id)).unwrap(); - let index2 = state.alloc(Some(id)).unwrap(); + let index1 = state.alloc(Some((id, memory_index))).unwrap(); + let index2 = state.alloc(Some((id, memory_index))).unwrap(); state.free(index2); state.free(index1); - assert!(state.alloc_affine_and_clear_affinity(id).is_some()); - assert!(state.alloc_affine_and_clear_affinity(id).is_some()); - assert_eq!(state.alloc_affine_and_clear_affinity(id), None); + assert!(state + .alloc_affine_and_clear_affinity(id, memory_index) + .is_some()); + assert!(state + .alloc_affine_and_clear_affinity(id, memory_index) + .is_some()); + assert_eq!( + state.alloc_affine_and_clear_affinity(id, memory_index), + None + ); } } @@ -525,7 +559,7 @@ mod test { let mut rng = rand::thread_rng(); let id_alloc = CompiledModuleIdAllocator::new(); - let ids = std::iter::repeat_with(|| id_alloc.alloc()) + let ids = std::iter::repeat_with(|| (id_alloc.alloc(), DefinedMemoryIndex::new(0))) .take(10) .collect::>(); let state = ModuleAffinityIndexAllocator::new(1000, 1000); @@ -569,9 +603,9 @@ mod test { #[test] fn test_affinity_threshold() { let id_alloc = CompiledModuleIdAllocator::new(); - let id1 = id_alloc.alloc(); - let id2 = id_alloc.alloc(); - let id3 = id_alloc.alloc(); + let id1 = (id_alloc.alloc(), DefinedMemoryIndex::new(0)); + let id2 = (id_alloc.alloc(), DefinedMemoryIndex::new(0)); + let id3 = (id_alloc.alloc(), DefinedMemoryIndex::new(0)); let state = ModuleAffinityIndexAllocator::new(10, 2); // Set some slot affinities @@ -605,13 +639,22 @@ mod test { state.free(SlotId(0)); // LRU is 1, so that should be picked - assert_eq!(state.alloc(Some(id_alloc.alloc())), Some(SlotId(1))); + assert_eq!( + state.alloc(Some((id_alloc.alloc(), DefinedMemoryIndex::new(0)))), + Some(SlotId(1)) + ); // Pick another LRU entry, this time 2 - assert_eq!(state.alloc(Some(id_alloc.alloc())), Some(SlotId(2))); + assert_eq!( + state.alloc(Some((id_alloc.alloc(), DefinedMemoryIndex::new(0)))), + Some(SlotId(2)) + ); // This should preserve slot `0` and pick up something new - assert_eq!(state.alloc(Some(id_alloc.alloc())), Some(SlotId(3))); + assert_eq!( + state.alloc(Some((id_alloc.alloc(), DefinedMemoryIndex::new(0)))), + Some(SlotId(3)) + ); state.free(SlotId(1)); state.free(SlotId(2)); diff --git a/crates/runtime/src/instance/allocator/pooling/memory_pool.rs b/crates/runtime/src/instance/allocator/pooling/memory_pool.rs new file mode 100644 index 000000000000..ba72e05b90c2 --- /dev/null +++ b/crates/runtime/src/instance/allocator/pooling/memory_pool.rs @@ -0,0 +1,439 @@ +use super::{ + index_allocator::{ModuleAffinityIndexAllocator, SlotId}, + MemoryAllocationIndex, +}; +use crate::{ + CompiledModuleId, InstanceAllocationRequest, Memory, MemoryImageSlot, Mmap, + PoolingInstanceAllocatorConfig, +}; +use anyhow::{anyhow, bail, Context, Result}; +use libc::c_void; +use std::sync::Mutex; +use wasmtime_environ::{ + DefinedMemoryIndex, MemoryPlan, MemoryStyle, Module, Tunables, WASM_PAGE_SIZE, +}; + +/// Represents a pool of WebAssembly linear memories. +/// +/// A linear memory is divided into accessible pages and guard pages. +/// +/// Each instance index into the pool returns an iterator over the base +/// addresses of the instance's linear memories. +/// +/// A diagram for this struct's fields is: +/// +/// ```ignore +/// memory_size +/// / +/// max_accessible / memory_and_guard_size +/// | / | +/// <--+---> / <-----------+----------> +/// <--------+-> +/// +/// +-----------+--------+---+-----------+ +--------+---+-----------+ +/// | PROT_NONE | | PROT_NONE | ... | | PROT_NONE | +/// +-----------+--------+---+-----------+ +--------+---+-----------+ +/// | |<------------------+----------------------------------> +/// \ | \ +/// mapping | `max_instances * max_memories` memories +/// / +/// initial_memory_offset +/// ``` +#[derive(Debug)] +pub struct MemoryPool { + mapping: Mmap, + index_allocator: ModuleAffinityIndexAllocator, + // If using a copy-on-write allocation scheme, the slot management. We + // dynamically transfer ownership of a slot to a Memory when in + // use. + image_slots: Vec>>, + // The size, in bytes, of each linear memory's reservation, not including + // any guard region. + memory_size: usize, + // The size, in bytes, of each linear memory's reservation plus the trailing + // guard region allocated for it. + memory_and_guard_size: usize, + // The maximum size that can become accessible, in bytes, of each linear + // memory. Guaranteed to be a whole number of wasm pages. + max_accessible: usize, + // The size, in bytes, of the offset to the first linear memory in this + // pool. This is here to help account for the first region of guard pages, + // if desired, before the first linear memory. + initial_memory_offset: usize, + // The maximum number of memories that can be allocated concurrently, aka + // our pool's capacity. + max_total_memories: usize, + // The maximum number of memories that a single core module instance may + // use. + memories_per_instance: usize, + // How much linear memory, in bytes, to keep resident after resetting for + // use with the next instance. This much memory will be `memset` to zero + // when a linear memory is deallocated. + // + // Memory exceeding this amount in the wasm linear memory will be released + // with `madvise` back to the kernel. + // + // Only applicable on Linux. + keep_resident: usize, +} + +impl MemoryPool { + /// Create a new `MemoryPool`. + pub fn new(config: &PoolingInstanceAllocatorConfig, tunables: &Tunables) -> Result { + // The maximum module memory page count cannot exceed 65536 pages + if config.limits.memory_pages > 0x10000 { + bail!( + "module memory page limit of {} exceeds the maximum of 65536", + config.limits.memory_pages + ); + } + + // Interpret the larger of the maximal size of memory or the static + // memory bound as the size of the virtual address space reservation for + // memory itself. Typically `static_memory_bound` is 4G which helps + // elide most bounds checks in wasm. If `memory_pages` is larger, + // though, then this is a non-moving pooling allocator so create larger + // reservations for account for that. + let memory_size = config.limits.memory_pages.max(tunables.static_memory_bound) + * u64::from(WASM_PAGE_SIZE); + + let memory_and_guard_size = + usize::try_from(memory_size + tunables.static_memory_offset_guard_size) + .map_err(|_| anyhow!("memory reservation size exceeds addressable memory"))?; + + assert!( + memory_and_guard_size % crate::page_size() == 0, + "memory size {} is not a multiple of system page size", + memory_and_guard_size + ); + + let max_total_memories = config.limits.total_memories as usize; + let initial_memory_offset = if tunables.guard_before_linear_memory { + usize::try_from(tunables.static_memory_offset_guard_size).unwrap() + } else { + 0 + }; + + // The entire allocation here is the size of each memory times the + // max memories per instance times the number of instances allowed in + // this pool, plus guard regions. + // + // Note, though, that guard regions are required to be after each linear + // memory. If the `guard_before_linear_memory` setting is specified, + // then due to the contiguous layout of linear memories the guard pages + // after one memory are also guard pages preceding the next linear + // memory. This means that we only need to handle pre-guard-page sizes + // specially for the first linear memory, hence the + // `initial_memory_offset` variable here. If guards aren't specified + // before linear memories this is set to `0`, otherwise it's set to + // the same size as guard regions for other memories. + let allocation_size = memory_and_guard_size + .checked_mul(max_total_memories) + .and_then(|c| c.checked_add(initial_memory_offset)) + .ok_or_else(|| { + anyhow!("total size of memory reservation exceeds addressable memory") + })?; + + // Create a completely inaccessible region to start + let mapping = Mmap::accessible_reserved(0, allocation_size) + .context("failed to create memory pool mapping")?; + + let image_slots: Vec<_> = std::iter::repeat_with(|| Mutex::new(None)) + .take(max_total_memories) + .collect(); + + let pool = Self { + index_allocator: ModuleAffinityIndexAllocator::new( + config.limits.total_memories, + config.max_unused_warm_slots, + ), + mapping, + image_slots, + memory_size: memory_size.try_into().unwrap(), + memory_and_guard_size, + initial_memory_offset, + max_total_memories, + memories_per_instance: usize::try_from(config.limits.max_memories_per_module).unwrap(), + max_accessible: (config.limits.memory_pages as usize) * (WASM_PAGE_SIZE as usize), + keep_resident: config.linear_memory_keep_resident, + }; + + Ok(pool) + } + + /// Validate whether this memory pool supports the given module. + pub fn validate(&self, module: &Module) -> Result<()> { + let memories = module.memory_plans.len() - module.num_imported_memories; + if memories > usize::try_from(self.memories_per_instance).unwrap() { + bail!( + "defined memories count of {} exceeds the per-instance limit of {}", + memories, + self.memories_per_instance, + ); + } + + for (i, plan) in module + .memory_plans + .iter() + .skip(module.num_imported_memories) + { + match plan.style { + MemoryStyle::Static { bound } => { + if u64::try_from(self.memory_size).unwrap() < bound { + bail!( + "memory size allocated per-memory is too small to \ + satisfy static bound of {bound:#x} pages" + ); + } + } + MemoryStyle::Dynamic { .. } => {} + } + let max = self.max_accessible / (WASM_PAGE_SIZE as usize); + if plan.memory.minimum > u64::try_from(max).unwrap() { + bail!( + "memory index {} has a minimum page size of {} which exceeds the limit of {}", + i.as_u32(), + plan.memory.minimum, + max, + ); + } + } + Ok(()) + } + + /// Allocate a single memory for the given instance allocation request. + pub fn allocate( + &self, + request: &mut InstanceAllocationRequest, + memory_plan: &MemoryPlan, + memory_index: DefinedMemoryIndex, + ) -> Result<(MemoryAllocationIndex, Memory)> { + let allocation_index = self + .index_allocator + .alloc( + request + .runtime_info + .unique_id() + .map(|id| (id, memory_index)), + ) + .map(|slot| MemoryAllocationIndex(u32::try_from(slot.index()).unwrap())) + .ok_or_else(|| { + anyhow!( + "maximum concurrent memory limit of {} reached", + self.max_total_memories + ) + })?; + + // Double-check that the runtime requirements of the memory are + // satisfied by the configuration of this pooling allocator. This + // should be returned as an error through `validate_memory_plans` + // but double-check here to be sure. + match memory_plan.style { + MemoryStyle::Static { bound } => { + let bound = bound * u64::from(WASM_PAGE_SIZE); + assert!(bound <= u64::try_from(self.memory_size).unwrap()); + } + MemoryStyle::Dynamic { .. } => {} + } + + let base_ptr = self.get_base(allocation_index); + let base_capacity = self.max_accessible; + + let mut slot = self.take_memory_image_slot(allocation_index); + let image = request.runtime_info.memory_image(memory_index)?; + let initial_size = memory_plan.memory.minimum * WASM_PAGE_SIZE as u64; + + // If instantiation fails, we can propagate the error + // upward and drop the slot. This will cause the Drop + // handler to attempt to map the range with PROT_NONE + // memory, to reserve the space while releasing any + // stale mappings. The next use of this slot will then + // create a new slot that will try to map over + // this, returning errors as well if the mapping + // errors persist. The unmap-on-drop is best effort; + // if it fails, then we can still soundly continue + // using the rest of the pool and allowing the rest of + // the process to continue, because we never perform a + // mmap that would leave an open space for someone + // else to come in and map something. + slot.instantiate(initial_size as usize, image, memory_plan)?; + + let memory = Memory::new_static( + memory_plan, + base_ptr, + base_capacity, + slot, + self.memory_and_guard_size, + unsafe { &mut *request.store.get().unwrap() }, + )?; + + Ok((allocation_index, memory)) + } + + /// Deallocate a previously-allocated memory. + /// + /// # Safety + /// + /// The memory must have been previously allocated from this pool and + /// assigned the given index, must currently be in an allocated state, and + /// must never be used again. + pub unsafe fn deallocate(&self, allocation_index: MemoryAllocationIndex, memory: Memory) { + let mut image = memory.unwrap_static_image(); + + // 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 + // reservation. + if image.clear_and_remain_ready(self.keep_resident).is_ok() { + self.return_memory_image_slot(allocation_index, image); + } + + self.index_allocator.free(SlotId(allocation_index.0)); + } + + /// Purging everything related to `module`. + pub fn purge_module(&self, module: CompiledModuleId) { + // This primarily means clearing out all of its memory images present in + // the virtual address space. Go through the index allocator for slots + // affine to `module` and reset them, freeing up the index when we're + // done. + // + // Note that this is only called when the specified `module` won't be + // allocated further (the module is being dropped) so this shouldn't hit + // any sort of infinite loop since this should be the final operation + // working with `module`. + for i in 0..self.memories_per_instance { + use wasmtime_environ::EntityRef; + let memory_index = DefinedMemoryIndex::new(i); + while let Some(id) = self + .index_allocator + .alloc_affine_and_clear_affinity(module, memory_index) + { + // Clear the image from the slot and, if successful, return it back + // to our state. Note that on failure here the whole slot will get + // paved over with an anonymous mapping. + let index = MemoryAllocationIndex(id.0); + let mut slot = self.take_memory_image_slot(index); + if slot.remove_image().is_ok() { + self.return_memory_image_slot(index, slot); + } + + self.index_allocator.free(id); + } + } + } + + fn get_base(&self, allocation_index: MemoryAllocationIndex) -> *mut u8 { + assert!(allocation_index.index() < self.max_total_memories); + let offset = + self.initial_memory_offset + allocation_index.index() * self.memory_and_guard_size; + unsafe { self.mapping.as_ptr().offset(offset as isize).cast_mut() } + } + + /// Take ownership of the given image slot. Must be returned via + /// `return_memory_image_slot` when the instance is done using it. + fn take_memory_image_slot(&self, allocation_index: MemoryAllocationIndex) -> MemoryImageSlot { + let maybe_slot = self.image_slots[allocation_index.index()] + .lock() + .unwrap() + .take(); + + maybe_slot.unwrap_or_else(|| { + MemoryImageSlot::create( + self.get_base(allocation_index) as *mut c_void, + 0, + self.max_accessible, + ) + }) + } + + /// Return ownership of the given image slot. + fn return_memory_image_slot( + &self, + allocation_index: MemoryAllocationIndex, + slot: MemoryImageSlot, + ) { + assert!(!slot.is_dirty()); + *self.image_slots[allocation_index.index()].lock().unwrap() = Some(slot); + } +} + +impl Drop for MemoryPool { + fn drop(&mut self) { + // Clear the `clear_no_drop` flag (i.e., ask to *not* clear on + // drop) for all slots, and then drop them here. This is + // valid because the one `Mmap` that covers the whole region + // can just do its one munmap. + for mut slot in std::mem::take(&mut self.image_slots) { + if let Some(slot) = slot.get_mut().unwrap() { + slot.no_clear_on_drop(); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{InstanceLimits, PoolingInstanceAllocator}; + use wasmtime_environ::WASM_PAGE_SIZE; + + #[cfg(target_pointer_width = "64")] + #[test] + fn test_memory_pool() -> Result<()> { + let pool = MemoryPool::new( + &PoolingInstanceAllocatorConfig { + limits: InstanceLimits { + total_memories: 5, + max_tables_per_module: 0, + max_memories_per_module: 3, + table_elements: 0, + memory_pages: 1, + ..Default::default() + }, + ..Default::default() + }, + &Tunables { + static_memory_bound: 1, + static_memory_offset_guard_size: 0, + ..Tunables::default() + }, + )?; + + assert_eq!(pool.memory_and_guard_size, WASM_PAGE_SIZE as usize); + assert_eq!(pool.max_total_memories, 5); + assert_eq!(pool.max_accessible, WASM_PAGE_SIZE as usize); + + let base = pool.mapping.as_ptr() as usize; + + for i in 0..5 { + let index = MemoryAllocationIndex(i); + let ptr = pool.get_base(index); + assert_eq!(ptr as usize - base, i as usize * pool.memory_and_guard_size); + } + + Ok(()) + } + + #[test] + fn test_pooling_allocator_with_reservation_size_exceeded() { + let config = PoolingInstanceAllocatorConfig { + limits: InstanceLimits { + total_memories: 1, + memory_pages: 2, + ..Default::default() + }, + ..PoolingInstanceAllocatorConfig::default() + }; + let pool = PoolingInstanceAllocator::new( + &config, + &Tunables { + static_memory_bound: 1, + static_memory_offset_guard_size: 0, + ..Tunables::default() + }, + ) + .unwrap(); + assert_eq!(pool.memories.memory_size, 2 * 65536); + } +} diff --git a/crates/runtime/src/instance/allocator/pooling/simple.rs b/crates/runtime/src/instance/allocator/pooling/simple.rs new file mode 100644 index 000000000000..f818d749b69c --- /dev/null +++ b/crates/runtime/src/instance/allocator/pooling/simple.rs @@ -0,0 +1,24 @@ +//! A simple index allocator. +//! +//! This index allocator doesn't do any module affinity or anything like that, +//! however it is built on top of the `ModuleAffinityIndexAllocator` to save +//! code (and code size). + +use super::module_affinity::{ModuleAffinityIndexAllocator, SlotId}; + +#[derive(Debug)] +pub struct SimpleIndexAllocator(ModuleAffinityIndexAllocator); + +impl SimpleIndexAllocator { + pub fn new(max_instances: u32) -> Self { + SimpleIndexAllocator(ModuleAffinityIndexAllocator::new(max_instances, 0)) + } + + pub fn alloc(&self) -> Option { + self.0.alloc(None) + } + + pub(crate) fn free(&self, index: SlotId) { + self.0.free(index); + } +} diff --git a/crates/runtime/src/instance/allocator/pooling/stack_pool.rs b/crates/runtime/src/instance/allocator/pooling/stack_pool.rs new file mode 100644 index 000000000000..651e212646a7 --- /dev/null +++ b/crates/runtime/src/instance/allocator/pooling/stack_pool.rs @@ -0,0 +1,237 @@ +use super::{ + imp::{commit_stack_pages, reset_stack_pages_to_zero}, + index_allocator::{SimpleIndexAllocator, SlotId}, + round_up_to_pow2, +}; +use crate::{Mmap, PoolingInstanceAllocatorConfig}; +use anyhow::{anyhow, bail, Context, Result}; + +/// Represents a pool of execution stacks (used for the async fiber implementation). +/// +/// Each index into the pool represents a single execution stack. The maximum number of +/// stacks is the same as the maximum number of instances. +/// +/// As stacks grow downwards, each stack starts (lowest address) with a guard page +/// that can be used to detect stack overflow. +/// +/// The top of the stack (starting stack pointer) is returned when a stack is allocated +/// from the pool. +#[derive(Debug)] +pub struct StackPool { + mapping: Mmap, + stack_size: usize, + max_stacks: usize, + page_size: usize, + index_allocator: SimpleIndexAllocator, + async_stack_zeroing: bool, + async_stack_keep_resident: usize, +} + +impl StackPool { + pub fn new(config: &PoolingInstanceAllocatorConfig) -> Result { + use rustix::mm::{mprotect, MprotectFlags}; + + let page_size = crate::page_size(); + + // Add a page to the stack size for the guard page when using fiber stacks + let stack_size = if config.stack_size == 0 { + 0 + } else { + round_up_to_pow2(config.stack_size, page_size) + .checked_add(page_size) + .ok_or_else(|| anyhow!("stack size exceeds addressable memory"))? + }; + + let max_stacks = usize::try_from(config.limits.total_stacks).unwrap(); + + let allocation_size = stack_size + .checked_mul(max_stacks) + .ok_or_else(|| anyhow!("total size of execution stacks exceeds addressable memory"))?; + + let mapping = Mmap::accessible_reserved(allocation_size, allocation_size) + .context("failed to create stack pool mapping")?; + + // Set up the stack guard pages. + if allocation_size > 0 { + unsafe { + for i in 0..max_stacks { + // Make the stack guard page inaccessible. + let bottom_of_stack = mapping.as_ptr().add(i * stack_size).cast_mut(); + mprotect(bottom_of_stack.cast(), page_size, MprotectFlags::empty()) + .context("failed to protect stack guard page")?; + } + } + } + + Ok(Self { + mapping, + stack_size, + max_stacks, + page_size, + async_stack_zeroing: config.async_stack_zeroing, + async_stack_keep_resident: config.async_stack_keep_resident, + index_allocator: SimpleIndexAllocator::new(config.limits.total_stacks), + }) + } + + /// Allocate a new fiber. + pub fn allocate(&self) -> Result { + if self.stack_size == 0 { + bail!("pooling allocator not configured to enable fiber stack allocation"); + } + + let index = self + .index_allocator + .alloc() + .ok_or_else(|| { + anyhow!( + "maximum concurrent fiber limit of {} reached", + self.max_stacks + ) + })? + .index(); + + assert!(index < self.max_stacks); + + unsafe { + // Remove the guard page from the size + let size_without_guard = self.stack_size - self.page_size; + + let bottom_of_stack = self + .mapping + .as_ptr() + .add((index * self.stack_size) + self.page_size) + .cast_mut(); + + commit_stack_pages(bottom_of_stack, size_without_guard)?; + + let stack = + wasmtime_fiber::FiberStack::from_raw_parts(bottom_of_stack, size_without_guard)?; + Ok(stack) + } + } + + /// Deallocate a previously-allocated fiber. + /// + /// # Safety + /// + /// The fiber must have been allocated by this pool, must be in an allocated + /// state, and must never be used again. + pub unsafe fn deallocate(&self, stack: &wasmtime_fiber::FiberStack) { + let top = stack + .top() + .expect("fiber stack not allocated from the pool") as usize; + + let base = self.mapping.as_ptr() as usize; + let len = self.mapping.len(); + assert!( + top > base && top <= (base + len), + "fiber stack top pointer not in range" + ); + + // Remove the guard page from the size + let stack_size = self.stack_size - self.page_size; + let bottom_of_stack = top - stack_size; + let start_of_stack = bottom_of_stack - self.page_size; + assert!(start_of_stack >= base && start_of_stack < (base + len)); + assert!((start_of_stack - base) % self.stack_size == 0); + + let index = (start_of_stack - base) / self.stack_size; + assert!(index < self.max_stacks); + + if self.async_stack_zeroing { + self.zero_stack(bottom_of_stack, stack_size); + } + + self.index_allocator.free(SlotId(index as u32)); + } + + fn zero_stack(&self, bottom: usize, size: usize) { + // Manually zero the top of the stack to keep the pages resident in + // memory and avoid future page faults. Use the system to deallocate + // pages past this. This hopefully strikes a reasonable balance between: + // + // * memset for the whole range is probably expensive + // * madvise for the whole range incurs expensive future page faults + // * most threads probably don't use most of the stack anyway + let size_to_memset = size.min(self.async_stack_keep_resident); + unsafe { + std::ptr::write_bytes( + (bottom + size - size_to_memset) as *mut u8, + 0, + size_to_memset, + ); + } + + // Use the system to reset remaining stack pages to zero. + reset_stack_pages_to_zero(bottom as _, size - size_to_memset).unwrap(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::InstanceLimits; + + #[cfg(all(unix, target_pointer_width = "64", feature = "async", not(miri)))] + #[test] + fn test_stack_pool() -> Result<()> { + let config = PoolingInstanceAllocatorConfig { + limits: InstanceLimits { + total_stacks: 10, + ..Default::default() + }, + stack_size: 1, + async_stack_zeroing: true, + ..PoolingInstanceAllocatorConfig::default() + }; + let pool = StackPool::new(&config)?; + + let native_page_size = crate::page_size(); + assert_eq!(pool.stack_size, 2 * native_page_size); + assert_eq!(pool.max_stacks, 10); + assert_eq!(pool.page_size, native_page_size); + + assert_eq!(pool.index_allocator.testing_freelist(), []); + + let base = pool.mapping.as_ptr() as usize; + + let mut stacks = Vec::new(); + for i in 0..10 { + let stack = pool.allocate().expect("allocation should succeed"); + assert_eq!( + ((stack.top().unwrap() as usize - base) / pool.stack_size) - 1, + i + ); + stacks.push(stack); + } + + assert_eq!(pool.index_allocator.testing_freelist(), []); + + pool.allocate().unwrap_err(); + + for stack in stacks { + unsafe { + pool.deallocate(&stack); + } + } + + assert_eq!( + pool.index_allocator.testing_freelist(), + [ + SlotId(0), + SlotId(1), + SlotId(2), + SlotId(3), + SlotId(4), + SlotId(5), + SlotId(6), + SlotId(7), + SlotId(8), + SlotId(9) + ], + ); + + Ok(()) + } +} diff --git a/crates/runtime/src/instance/allocator/pooling/table_pool.rs b/crates/runtime/src/instance/allocator/pooling/table_pool.rs new file mode 100644 index 000000000000..a5607c46c4de --- /dev/null +++ b/crates/runtime/src/instance/allocator/pooling/table_pool.rs @@ -0,0 +1,210 @@ +use super::{ + imp::{commit_table_pages, decommit_table_pages}, + index_allocator::{SimpleIndexAllocator, SlotId}, + round_up_to_pow2, TableAllocationIndex, +}; +use crate::{InstanceAllocationRequest, Mmap, PoolingInstanceAllocatorConfig, Table}; +use anyhow::{anyhow, bail, Context, Result}; +use std::mem; +use wasmtime_environ::{Module, TablePlan}; + +/// Represents a pool of WebAssembly tables. +/// +/// Each instance index into the pool returns an iterator over the base addresses +/// of the instance's tables. +#[derive(Debug)] +pub struct TablePool { + index_allocator: SimpleIndexAllocator, + mapping: Mmap, + table_size: usize, + max_total_tables: usize, + tables_per_instance: usize, + page_size: usize, + keep_resident: usize, + table_elements: usize, +} + +impl TablePool { + /// Create a new `TablePool`. + pub fn new(config: &PoolingInstanceAllocatorConfig) -> Result { + let page_size = crate::page_size(); + + let table_size = round_up_to_pow2( + mem::size_of::<*mut u8>() + .checked_mul(config.limits.table_elements as usize) + .ok_or_else(|| anyhow!("table size exceeds addressable memory"))?, + page_size, + ); + + let max_total_tables = usize::try_from(config.limits.total_tables).unwrap(); + let tables_per_instance = usize::try_from(config.limits.max_tables_per_module).unwrap(); + + let allocation_size = table_size + .checked_mul(max_total_tables) + .ok_or_else(|| anyhow!("total size of tables exceeds addressable memory"))?; + + let mapping = Mmap::accessible_reserved(allocation_size, allocation_size) + .context("failed to create table pool mapping")?; + + Ok(Self { + index_allocator: SimpleIndexAllocator::new(config.limits.total_tables), + mapping, + table_size, + max_total_tables, + tables_per_instance, + page_size, + keep_resident: config.table_keep_resident, + table_elements: usize::try_from(config.limits.table_elements).unwrap(), + }) + } + + /// Validate whether this module's tables are allocatable by this pool. + pub fn validate(&self, module: &Module) -> Result<()> { + let tables = module.table_plans.len() - module.num_imported_tables; + + if tables > usize::try_from(self.tables_per_instance).unwrap() { + bail!( + "defined tables count of {} exceeds the per-instance limit of {}", + tables, + self.tables_per_instance, + ); + } + + if tables > self.max_total_tables { + bail!( + "defined tables count of {} exceeds the total tables limit of {}", + tables, + self.max_total_tables, + ); + } + + for (i, plan) in module.table_plans.iter().skip(module.num_imported_tables) { + if plan.table.minimum > u32::try_from(self.table_elements).unwrap() { + bail!( + "table index {} has a minimum element size of {} which exceeds the limit of {}", + i.as_u32(), + plan.table.minimum, + self.table_elements, + ); + } + } + Ok(()) + } + + /// Get the base pointer of the given table allocation. + fn get(&self, table_index: TableAllocationIndex) -> *mut u8 { + assert!(table_index.index() < self.max_total_tables); + + unsafe { + self.mapping + .as_ptr() + .add(table_index.index() * self.table_size) + .cast_mut() + } + } + + /// Allocate a single table for the given instance allocation request. + pub fn allocate( + &self, + request: &mut InstanceAllocationRequest, + table_plan: &TablePlan, + ) -> Result<(TableAllocationIndex, Table)> { + let allocation_index = self + .index_allocator + .alloc() + .map(|slot| TableAllocationIndex(slot.0)) + .ok_or_else(|| { + anyhow!( + "maximum concurrent table limit of {} reached", + self.max_total_tables + ) + })?; + + let base = self.get(allocation_index); + + commit_table_pages( + base as *mut u8, + self.table_elements * mem::size_of::<*mut u8>(), + )?; + + let table = Table::new_static( + table_plan, + unsafe { std::slice::from_raw_parts_mut(base.cast(), self.table_elements) }, + unsafe { &mut *request.store.get().unwrap() }, + )?; + + Ok((allocation_index, table)) + } + + /// Deallocate a previously-allocated table. + /// + /// # Safety + /// + /// The table must have been previously-allocated by this pool and assigned + /// the given allocation index, it must currently be allocated, and it must + /// never be used again. + pub unsafe fn deallocate(&self, allocation_index: TableAllocationIndex, table: Table) { + assert!(table.is_static()); + + let size = round_up_to_pow2( + table.size() as usize * mem::size_of::<*mut u8>(), + self.page_size, + ); + + drop(table); + + let base = self.get(allocation_index); + self.reset_table_pages_to_zero(base, size) + .expect("failed to decommit table pages"); + + self.index_allocator.free(SlotId(allocation_index.0)); + } + + fn reset_table_pages_to_zero(&self, base: *mut u8, size: usize) -> Result<()> { + let size_to_memset = size.min(self.keep_resident); + unsafe { + std::ptr::write_bytes(base, 0, size_to_memset); + decommit_table_pages(base.add(size_to_memset), size - size_to_memset) + .context("failed to decommit table page")?; + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::InstanceLimits; + + #[cfg(target_pointer_width = "64")] + #[test] + fn test_table_pool() -> Result<()> { + let pool = TablePool::new(&PoolingInstanceAllocatorConfig { + limits: InstanceLimits { + total_tables: 7, + table_elements: 100, + memory_pages: 0, + max_memories_per_module: 0, + ..Default::default() + }, + ..Default::default() + })?; + + let host_page_size = crate::page_size(); + + assert_eq!(pool.table_size, host_page_size); + assert_eq!(pool.max_total_tables, 7); + assert_eq!(pool.page_size, host_page_size); + assert_eq!(pool.table_elements, 100); + + let base = pool.mapping.as_ptr() as usize; + + for i in 0..7 { + let index = TableAllocationIndex(i); + let ptr = pool.get(index); + assert_eq!(ptr as usize - base, i as usize * pool.table_size); + } + + Ok(()) + } +} diff --git a/crates/runtime/src/lib.rs b/crates/runtime/src/lib.rs index 23896612af6e..ce4a2d576aa2 100644 --- a/crates/runtime/src/lib.rs +++ b/crates/runtime/src/lib.rs @@ -55,8 +55,8 @@ pub use crate::export::*; pub use crate::externref::*; pub use crate::imports::Imports; pub use crate::instance::{ - Instance, InstanceAllocationRequest, InstanceAllocator, InstanceHandle, - OnDemandInstanceAllocator, StorePtr, + Instance, InstanceAllocationRequest, InstanceAllocator, InstanceAllocatorImpl, InstanceHandle, + MemoryAllocationIndex, OnDemandInstanceAllocator, StorePtr, TableAllocationIndex, }; #[cfg(feature = "pooling-allocator")] pub use crate::instance::{ diff --git a/crates/wasi/src/preview2/preview1.rs b/crates/wasi/src/preview2/preview1.rs index a2f957aa1804..be46d2284ea5 100644 --- a/crates/wasi/src/preview2/preview1.rs +++ b/crates/wasi/src/preview2/preview1.rs @@ -1190,7 +1190,7 @@ impl< .. }) if self.table().is_file(fd) => { let Some(buf) = first_non_empty_iovec(iovs)? else { - return Ok(0) + return Ok(0); }; let pos = position.load(Ordering::Relaxed); @@ -1215,7 +1215,7 @@ impl< } Descriptor::Stdin { input_stream, .. } => { let Some(buf) = first_non_empty_iovec(iovs)? else { - return Ok(0) + return Ok(0); }; let (read, state) = streams::Host::read( self, @@ -1253,7 +1253,7 @@ impl< let (mut buf, read, state) = match desc { Descriptor::File(File { fd, blocking, .. }) if self.table().is_file(fd) => { let Some(buf) = first_non_empty_iovec(iovs)? else { - return Ok(0) + return Ok(0); }; let stream = self.read_via_stream(fd, offset).await.map_err(|e| { @@ -1306,7 +1306,7 @@ impl< position, }) if self.table().is_file(fd) => { let Some(buf) = first_non_empty_ciovec(ciovs)? else { - return Ok(0) + return Ok(0); }; let (stream, pos) = if append { let stream = self.append_via_stream(fd).await.map_err(|e| { @@ -1338,7 +1338,7 @@ impl< } Descriptor::Stdout { output_stream, .. } | Descriptor::Stderr { output_stream, .. } => { let Some(buf) = first_non_empty_ciovec(ciovs)? else { - return Ok(0) + return Ok(0); }; let (n, _stat) = streams::Host::blocking_write(self, output_stream, buf) .await @@ -1364,7 +1364,7 @@ impl< let (n, _stat) = match desc { Descriptor::File(File { fd, blocking, .. }) if self.table().is_file(fd) => { let Some(buf) = first_non_empty_ciovec(ciovs)? else { - return Ok(0) + return Ok(0); }; let stream = self.write_via_stream(fd, offset).await.map_err(|e| { e.try_into() diff --git a/crates/wasmtime/src/component/component.rs b/crates/wasmtime/src/component/component.rs index 45f4f007c514..6e4034a12e20 100644 --- a/crates/wasmtime/src/component/component.rs +++ b/crates/wasmtime/src/component/component.rs @@ -10,9 +10,9 @@ use std::ptr::NonNull; use std::sync::Arc; use wasmtime_environ::component::{ AllCallFunc, ComponentTypes, GlobalInitializer, InstantiateModule, StaticModuleIndex, - TrampolineIndex, Translator, + TrampolineIndex, Translator, VMComponentOffsets, }; -use wasmtime_environ::{FunctionLoc, ObjectKind, PrimaryMap, ScopeVec}; +use wasmtime_environ::{FunctionLoc, HostPtr, ObjectKind, PrimaryMap, ScopeVec}; use wasmtime_jit::{CodeMemory, CompiledModuleInfo}; use wasmtime_runtime::component::ComponentRuntimeInfo; use wasmtime_runtime::{ @@ -258,6 +258,14 @@ impl Component { let types = Arc::new(types); let code = Arc::new(CodeObject::new(code_memory, signatures, types.into())); + // Validate that the component can be used with the current instance + // allocator. + engine.allocator().validate_component( + &info.component, + &VMComponentOffsets::new(HostPtr, &info.component), + &|module_index| &static_modules[module_index].module, + )?; + // Convert all information about static core wasm modules into actual // `Module` instances by converting each `CompiledModuleInfo`, the // `types` type information, and the code memory to a runtime object. diff --git a/crates/wasmtime/src/component/instance.rs b/crates/wasmtime/src/component/instance.rs index 6a664889d904..a6ee266c2f2f 100644 --- a/crates/wasmtime/src/component/instance.rs +++ b/crates/wasmtime/src/component/instance.rs @@ -587,10 +587,22 @@ impl InstancePre { fn instantiate_impl(&self, mut store: impl AsContextMut) -> Result { let mut store = store.as_context_mut(); - let mut i = Instantiator::new(&self.component, store.0, &self.imports); - i.run(&mut store)?; - let data = Box::new(i.data); - Ok(Instance(store.0.store_data_mut().insert(Some(data)))) + store + .engine() + .allocator() + .increment_component_instance_count()?; + let mut instantiator = Instantiator::new(&self.component, store.0, &self.imports); + instantiator.run(&mut store).map_err(|e| { + store + .engine() + .allocator() + .decrement_component_instance_count(); + e + })?; + let data = Box::new(instantiator.data); + let instance = Instance(store.0.store_data_mut().insert(Some(data))); + store.0.push_component_instance(instance); + Ok(instance) } } diff --git a/crates/wasmtime/src/config.rs b/crates/wasmtime/src/config.rs index d86c96a82b09..9439baba3e84 100644 --- a/crates/wasmtime/src/config.rs +++ b/crates/wasmtime/src/config.rs @@ -1939,7 +1939,7 @@ impl PoolingAllocationConfig { /// This means that the total set of used slots in the pooling instance /// allocator can impact the overall RSS usage of a program. /// - /// The default value for this option is 100. + /// The default value for this option is `100`. pub fn max_unused_warm_slots(&mut self, max: u32) -> &mut Self { self.config.max_unused_warm_slots = max; self @@ -2022,47 +2022,180 @@ impl PoolingAllocationConfig { self } - /// The maximum number of concurrent instances supported (default is 1000). + /// The maximum number of concurrent component instances supported (default + /// is `1000`). + /// + /// This provides an upper-bound on the total size of component + /// metadata-related allocations, along with + /// [`PoolingAllocationConfig::component_instance_size`]. The upper bound is + /// + /// ```text + /// total_component_instances * component_instance_size + /// ``` + /// + /// where `component_instance_size` is rounded up to the size and alignment + /// of the internal representation of the metadata. + pub fn total_component_instances(&mut self, count: u32) -> &mut Self { + self.config.limits.total_component_instances = count; + self + } + + /// The maximum size, in bytes, allocated for a component instance's + /// `VMComponentContext` metadata. + /// + /// The [`wasmtime::component::Instance`][crate::component::Instance] type + /// has a static size but its internal `VMComponentContext` is dynamically + /// sized depending on the component being instantiated. This size limit + /// loosely correlates to the size of the component, taking into account + /// factors such as: + /// + /// * number of lifted and lowered functions, + /// * number of memories + /// * number of inner instances + /// * number of resources + /// + /// If the allocated size per instance is too small then instantiation of a + /// module will fail at runtime with an error indicating how many bytes were + /// needed. + /// + /// The default value for this is 1MiB. + /// + /// This provides an upper-bound on the total size of component + /// metadata-related allocations, along with + /// [`PoolingAllocationConfig::total_component_instances`]. The upper bound is + /// + /// ```text + /// total_component_instances * component_instance_size + /// ``` + /// + /// where `component_instance_size` is rounded up to the size and alignment + /// of the internal representation of the metadata. + pub fn component_instance_size(&mut self, size: usize) -> &mut Self { + self.config.limits.component_instance_size = size; + self + } + + /// The maximum number of core instances a single component may contain + /// (default is `20`). + /// + /// This method (along with + /// [`PoolingAllocationConfig::max_memories_per_component`], + /// [`PoolingAllocationConfig::max_tables_per_component`], and + /// [`PoolingAllocationConfig::component_instance_size`]) allows you to cap + /// the amount of resources a single component allocation consumes. + /// + /// If a component will instantiate more core instances than `count`, then + /// the component will fail to instantiate. + pub fn max_core_instances_per_component(&mut self, count: u32) -> &mut Self { + self.config.limits.max_core_instances_per_component = count; + self + } + + /// The maximum number of Wasm linear memories that a single component may + /// transitively contain (default is `20`). + /// + /// This method (along with + /// [`PoolingAllocationConfig::max_core_instances_per_component`], + /// [`PoolingAllocationConfig::max_tables_per_component`], and + /// [`PoolingAllocationConfig::component_instance_size`]) allows you to cap + /// the amount of resources a single component allocation consumes. + /// + /// If a component transitively contains more linear memories than `count`, + /// then the component will fail to instantiate. + pub fn max_memories_per_component(&mut self, count: u32) -> &mut Self { + self.config.limits.max_memories_per_component = count; + self + } + + /// The maximum number of tables that a single component may transitively + /// contain (default is `20`). + /// + /// This method (along with + /// [`PoolingAllocationConfig::max_core_instances_per_component`], + /// [`PoolingAllocationConfig::max_memories_per_component`], + /// [`PoolingAllocationConfig::component_instance_size`]) allows you to cap + /// the amount of resources a single component allocation consumes. + /// + /// If a component will transitively contains more tables than `count`, then + /// the component will fail to instantiate. + pub fn max_tables_per_component(&mut self, count: u32) -> &mut Self { + self.config.limits.max_tables_per_component = count; + self + } + + /// The maximum number of concurrent Wasm linear memories supported (default + /// is `1000`). /// /// This value has a direct impact on the amount of memory allocated by the pooling /// instance allocator. /// - /// The pooling instance allocator allocates three memory pools with sizes depending on this value: + /// The pooling instance allocator allocates a memory pool, where each entry + /// in the pool contains the reserved address space for each linear memory + /// supported by an instance. + /// + /// The memory pool will reserve a large quantity of host process address + /// space to elide the bounds checks required for correct WebAssembly memory + /// semantics. Even with 64-bit address spaces, the address space is limited + /// when dealing with a large number of linear memories. + /// + /// For example, on Linux x86_64, the userland address space limit is 128 + /// TiB. That might seem like a lot, but each linear memory will *reserve* 6 + /// GiB of space by default. + pub fn total_memories(&mut self, count: u32) -> &mut Self { + self.config.limits.total_memories = count; + self + } + + /// The maximum number of concurrent tables supported (default is `1000`). /// - /// * An instance pool, where each entry in the pool can store the runtime representation - /// of an instance, including a maximal `VMContext` structure. + /// This value has a direct impact on the amount of memory allocated by the + /// pooling instance allocator. /// - /// * A memory pool, where each entry in the pool contains the reserved address space for each - /// linear memory supported by an instance. + /// The pooling instance allocator allocates a table pool, where each entry + /// in the pool contains the space needed for each WebAssembly table + /// supported by an instance (see `table_elements` to control the size of + /// each table). + pub fn total_tables(&mut self, count: u32) -> &mut Self { + self.config.limits.total_tables = count; + self + } + + /// The maximum number of execution stacks allowed for asynchronous + /// execution, when enabled (default is `1000`). /// - /// * A table pool, where each entry in the pool contains the space needed for each WebAssembly table - /// supported by an instance (see `table_elements` to control the size of each table). + /// This value has a direct impact on the amount of memory allocated by the + /// pooling instance allocator. + #[cfg(feature = "async")] + pub fn total_stacks(&mut self, count: u32) -> &mut Self { + self.config.limits.total_stacks = count; + self + } + + /// The maximum number of concurrent core instances supported (default is + /// `1000`). /// - /// Additionally, this value will also control the maximum number of execution stacks allowed for - /// asynchronous execution (one per instance), when enabled. + /// This provides an upper-bound on the total size of core instance + /// metadata-related allocations, along with + /// [`PoolingAllocationConfig::core_instance_size`]. The upper bound is /// - /// The memory pool will reserve a large quantity of host process address space to elide the bounds - /// checks required for correct WebAssembly memory semantics. Even for 64-bit address spaces, the - /// address space is limited when dealing with a large number of supported instances. + /// ```text + /// total_core_instances * core_instance_size + /// ``` /// - /// For example, on Linux x86_64, the userland address space limit is 128 TiB. That might seem like a lot, - /// but each linear memory will *reserve* 6 GiB of space by default. Multiply that by the number of linear - /// memories each instance supports and then by the number of supported instances and it becomes apparent - /// that address space can be exhausted depending on the number of supported instances. - pub fn instance_count(&mut self, count: u32) -> &mut Self { - self.config.limits.count = count; + /// where `core_instance_size` is rounded up to the size and alignment of + /// the internal representation of the metadata. + pub fn total_core_instances(&mut self, count: u32) -> &mut Self { + self.config.limits.total_core_instances = count; self } - /// The maximum size, in bytes, allocated for an instance and its - /// `VMContext`. + /// The maximum size, in bytes, allocated for a core instance's `VMContext` + /// metadata. /// - /// This amount of space is pre-allocated for `count` number of instances - /// and is used to store the runtime `wasmtime_runtime::Instance` structure - /// along with its adjacent `VMContext` structure. The `Instance` type has a - /// static size but `VMContext` is dynamically sized depending on the module - /// being instantiated. This size limit loosely correlates to the size of - /// the wasm module, taking into account factors such as: + /// The [`Instance`] type has a static size but its `VMContext` metadata is + /// dynamically sized depending on the module being instantiated. This size + /// limit loosely correlates to the size of the Wasm module, taking into + /// account factors such as: /// /// * number of functions /// * number of globals @@ -2072,71 +2205,91 @@ impl PoolingAllocationConfig { /// /// If the allocated size per instance is too small then instantiation of a /// module will fail at runtime with an error indicating how many bytes were - /// needed. This amount of bytes are committed to memory per-instance when - /// a pooling allocator is created. + /// needed. + /// + /// The default value for this is 1MiB. + /// + /// This provides an upper-bound on the total size of core instance + /// metadata-related allocations, along with + /// [`PoolingAllocationConfig::total_core_instances`]. The upper bound is + /// + /// ```text + /// total_core_instances * core_instance_size + /// ``` /// - /// The default value for this is 1MB. - pub fn instance_size(&mut self, size: usize) -> &mut Self { - self.config.limits.size = size; + /// where `core_instance_size` is rounded up to the size and alignment of + /// the internal representation of the metadata. + pub fn core_instance_size(&mut self, size: usize) -> &mut Self { + self.config.limits.core_instance_size = size; self } - /// The maximum number of defined tables for a module (default is 1). + /// The maximum number of defined tables for a core module (default is `1`). /// - /// This value controls the capacity of the `VMTableDefinition` table in each instance's - /// `VMContext` structure. + /// This value controls the capacity of the `VMTableDefinition` table in + /// each instance's `VMContext` structure. /// - /// The allocated size of the table will be `tables * sizeof(VMTableDefinition)` for each - /// instance regardless of how many tables are defined by an instance's module. - pub fn instance_tables(&mut self, tables: u32) -> &mut Self { - self.config.limits.tables = tables; + /// The allocated size of the table will be `tables * + /// sizeof(VMTableDefinition)` for each instance regardless of how many + /// tables are defined by an instance's module. + pub fn max_tables_per_module(&mut self, tables: u32) -> &mut Self { + self.config.limits.max_tables_per_module = tables; self } - /// The maximum table elements for any table defined in a module (default is 10000). + /// The maximum table elements for any table defined in a module (default is + /// `10000`). /// - /// If a table's minimum element limit is greater than this value, the module will - /// fail to instantiate. + /// If a table's minimum element limit is greater than this value, the + /// module will fail to instantiate. /// - /// If a table's maximum element limit is unbounded or greater than this value, - /// the maximum will be `table_elements` for the purpose of any `table.grow` instruction. + /// If a table's maximum element limit is unbounded or greater than this + /// value, the maximum will be `table_elements` for the purpose of any + /// `table.grow` instruction. /// - /// This value is used to reserve the maximum space for each supported table; table elements - /// are pointer-sized in the Wasmtime runtime. Therefore, the space reserved for each instance - /// is `tables * table_elements * sizeof::<*const ()>`. - pub fn instance_table_elements(&mut self, elements: u32) -> &mut Self { + /// This value is used to reserve the maximum space for each supported + /// table; table elements are pointer-sized in the Wasmtime runtime. + /// Therefore, the space reserved for each instance is `tables * + /// table_elements * sizeof::<*const ()>`. + pub fn table_elements(&mut self, elements: u32) -> &mut Self { self.config.limits.table_elements = elements; self } - /// The maximum number of defined linear memories for a module (default is 1). + /// The maximum number of defined linear memories for a module (default is + /// `1`). /// - /// This value controls the capacity of the `VMMemoryDefinition` table in each instance's - /// `VMContext` structure. + /// This value controls the capacity of the `VMMemoryDefinition` table in + /// each core instance's `VMContext` structure. /// - /// The allocated size of the table will be `memories * sizeof(VMMemoryDefinition)` for each - /// instance regardless of how many memories are defined by an instance's module. - pub fn instance_memories(&mut self, memories: u32) -> &mut Self { - self.config.limits.memories = memories; + /// The allocated size of the table will be `memories * + /// sizeof(VMMemoryDefinition)` for each core instance regardless of how + /// many memories are defined by the core instance's module. + pub fn max_memories_per_module(&mut self, memories: u32) -> &mut Self { + self.config.limits.max_memories_per_module = memories; self } - /// The maximum number of pages for any linear memory defined in a module (default is 160). + /// The maximum number of Wasm pages for any linear memory defined in a + /// module (default is `160`). /// - /// The default of 160 means at most 10 MiB of host memory may be committed for each instance. + /// The default of `160` means at most 10 MiB of host memory may be + /// committed for each instance. /// - /// If a memory's minimum page limit is greater than this value, the module will - /// fail to instantiate. + /// If a memory's minimum page limit is greater than this value, the module + /// will fail to instantiate. /// - /// If a memory's maximum page limit is unbounded or greater than this value, - /// the maximum will be `memory_pages` for the purpose of any `memory.grow` instruction. + /// If a memory's maximum page limit is unbounded or greater than this + /// value, the maximum will be `memory_pages` for the purpose of any + /// `memory.grow` instruction. /// - /// This value is used to control the maximum accessible space for each linear memory of an instance. + /// This value is used to control the maximum accessible space for each + /// linear memory of a core instance. /// /// The reservation size of each linear memory is controlled by the - /// `static_memory_maximum_size` setting and this value cannot - /// exceed the configured static memory maximum size. - pub fn instance_memory_pages(&mut self, pages: u64) -> &mut Self { + /// `static_memory_maximum_size` setting and this value cannot exceed the + /// configured static memory maximum size. + pub fn memory_pages(&mut self, pages: u64) -> &mut Self { self.config.limits.memory_pages = pages; self } diff --git a/crates/wasmtime/src/instance.rs b/crates/wasmtime/src/instance.rs index 13cde5749d91..1489903012ab 100644 --- a/crates/wasmtime/src/instance.rs +++ b/crates/wasmtime/src/instance.rs @@ -269,7 +269,7 @@ impl Instance { store .engine() .allocator() - .allocate(InstanceAllocationRequest { + .allocate_module(InstanceAllocationRequest { runtime_info: &module.runtime_info(), imports, host_state: Box::new(Instance(instance_to_be)), diff --git a/crates/wasmtime/src/module.rs b/crates/wasmtime/src/module.rs index 7d525d99c52e..3e812cfc8543 100644 --- a/crates/wasmtime/src/module.rs +++ b/crates/wasmtime/src/module.rs @@ -581,9 +581,11 @@ impl Module { engine.unique_id_allocator(), )?; - // Validate the module can be used with the current allocator + // Validate the module can be used with the current instance allocator. let offsets = VMOffsets::new(HostPtr, module.module()); - engine.allocator().validate(module.module(), &offsets)?; + engine + .allocator() + .validate_module(module.module(), &offsets)?; Ok(Self { inner: Arc::new(ModuleInner { diff --git a/crates/wasmtime/src/store.rs b/crates/wasmtime/src/store.rs index 788b69608a92..1df0e8c7ab28 100644 --- a/crates/wasmtime/src/store.rs +++ b/crates/wasmtime/src/store.rs @@ -291,6 +291,8 @@ pub struct StoreOpaque { engine: Engine, runtime_limits: VMRuntimeLimits, instances: Vec, + #[cfg(feature = "component-model")] + component_instances: Vec, signal_handler: Option>>, externref_activations_table: VMExternRefActivationsTable, modules: ModuleRegistry, @@ -461,6 +463,8 @@ impl Store { engine: engine.clone(), runtime_limits: Default::default(), instances: Vec::new(), + #[cfg(feature = "component-model")] + component_instances: Vec::new(), signal_handler: None, externref_activations_table: VMExternRefActivationsTable::new(), modules: ModuleRegistry::default(), @@ -505,15 +509,21 @@ impl Store { inner.default_caller = { let module = Arc::new(wasmtime_environ::Module::default()); let shim = BareModuleInfo::empty(module).into_traitobj(); - let mut instance = OnDemandInstanceAllocator::default() - .allocate(InstanceAllocationRequest { - host_state: Box::new(()), - imports: Default::default(), - store: StorePtr::empty(), - runtime_info: &shim, - wmemcheck: engine.config().wmemcheck, - }) - .expect("failed to allocate default callee"); + let allocator = OnDemandInstanceAllocator::default(); + allocator + .validate_module(shim.module(), shim.offsets()) + .unwrap(); + let mut instance = unsafe { + allocator + .allocate_module(InstanceAllocationRequest { + host_state: Box::new(()), + imports: Default::default(), + store: StorePtr::empty(), + runtime_info: &shim, + wmemcheck: engine.config().wmemcheck, + }) + .expect("failed to allocate default callee") + }; // Note the erasure of the lifetime here into `'static`, so in // general usage of this trait object must be strictly bounded to @@ -1579,6 +1589,11 @@ at https://bytecodealliance.org/security. ) { (&mut self.component_calls, &mut self.component_host_table) } + + #[cfg(feature = "component-model")] + pub(crate) fn push_component_instance(&mut self, instance: crate::component::Instance) { + self.component_instances.push(instance); + } } impl StoreContextMut<'_, T> { @@ -2181,12 +2196,19 @@ impl Drop for StoreOpaque { let ondemand = OnDemandInstanceAllocator::default(); for instance in self.instances.iter_mut() { if instance.ondemand { - ondemand.deallocate(&mut instance.handle); + ondemand.deallocate_module(&mut instance.handle); } else { - allocator.deallocate(&mut instance.handle); + allocator.deallocate_module(&mut instance.handle); + } + } + ondemand.deallocate_module(&mut self.default_caller); + + #[cfg(feature = "component-model")] + { + for _ in self.component_instances.iter() { + allocator.decrement_component_instance_count(); } } - ondemand.deallocate(&mut self.default_caller); // See documentation for these fields on `StoreOpaque` for why they // must be dropped in this order. diff --git a/crates/wasmtime/src/trampoline.rs b/crates/wasmtime/src/trampoline.rs index 0c7e45933632..928cc2cc2974 100644 --- a/crates/wasmtime/src/trampoline.rs +++ b/crates/wasmtime/src/trampoline.rs @@ -41,15 +41,14 @@ fn create_handle( let module = Arc::new(module); let runtime_info = &BareModuleInfo::maybe_imported_func(module, one_signature).into_traitobj(); - let handle = OnDemandInstanceAllocator::new(config.mem_creator.clone(), 0).allocate( - InstanceAllocationRequest { + let handle = OnDemandInstanceAllocator::new(config.mem_creator.clone(), 0) + .allocate_module(InstanceAllocationRequest { imports, host_state, store: StorePtr::new(store.traitobj()), runtime_info, wmemcheck: false, - }, - )?; + })?; Ok(store.add_instance(handle, true)) } diff --git a/crates/wasmtime/src/trampoline/memory.rs b/crates/wasmtime/src/trampoline/memory.rs index 3bd085d966de..a4ef32097960 100644 --- a/crates/wasmtime/src/trampoline/memory.rs +++ b/crates/wasmtime/src/trampoline/memory.rs @@ -7,13 +7,19 @@ use std::convert::TryFrom; use std::ops::Range; use std::sync::Arc; use wasmtime_environ::{ - DefinedMemoryIndex, DefinedTableIndex, EntityIndex, MemoryPlan, MemoryStyle, Module, - PrimaryMap, WASM_PAGE_SIZE, + DefinedMemoryIndex, DefinedTableIndex, EntityIndex, HostPtr, MemoryPlan, MemoryStyle, Module, + VMOffsets, WASM_PAGE_SIZE, }; use wasmtime_runtime::{ - CompiledModuleId, Imports, InstanceAllocationRequest, InstanceAllocator, Memory, MemoryImage, - OnDemandInstanceAllocator, RuntimeLinearMemory, RuntimeMemoryCreator, SharedMemory, StorePtr, - Table, VMMemoryDefinition, + CompiledModuleId, Imports, InstanceAllocationRequest, InstanceAllocator, InstanceAllocatorImpl, + Memory, MemoryAllocationIndex, MemoryImage, OnDemandInstanceAllocator, RuntimeLinearMemory, + RuntimeMemoryCreator, SharedMemory, StorePtr, Table, TableAllocationIndex, VMMemoryDefinition, +}; + +#[cfg(feature = "component-model")] +use wasmtime_environ::{ + component::{Component, VMComponentOffsets}, + StaticModuleIndex, }; /// Create a "frankenstein" instance with a single memory. @@ -64,7 +70,7 @@ pub fn create_memory( preallocation, ondemand: OnDemandInstanceAllocator::default(), } - .allocate(request)?; + .allocate_module(request)?; let instance_id = store.add_instance(handle.clone(), true); Ok(instance_id) } @@ -143,48 +149,94 @@ struct SingleMemoryInstance<'a> { ondemand: OnDemandInstanceAllocator, } -unsafe impl InstanceAllocator for SingleMemoryInstance<'_> { - fn allocate_index(&self, req: &InstanceAllocationRequest) -> Result { - self.ondemand.allocate_index(req) +unsafe impl InstanceAllocatorImpl for SingleMemoryInstance<'_> { + #[cfg(feature = "component-model")] + fn validate_component_impl<'a>( + &self, + _component: &Component, + _offsets: &VMComponentOffsets, + _get_module: &'a dyn Fn(StaticModuleIndex) -> &'a Module, + ) -> Result<()> { + unreachable!("`SingleMemoryInstance` allocator never used with components") + } + + fn validate_module_impl(&self, module: &Module, offsets: &VMOffsets) -> Result<()> { + anyhow::ensure!( + module.memory_plans.len() == 1, + "`SingleMemoryInstance` allocator can only be used for modules with a single memory" + ); + self.ondemand.validate_module_impl(module, offsets)?; + Ok(()) + } + + fn increment_component_instance_count(&self) -> Result<()> { + self.ondemand.increment_component_instance_count() + } + + fn decrement_component_instance_count(&self) { + self.ondemand.decrement_component_instance_count(); } - fn deallocate_index(&self, index: usize) { - self.ondemand.deallocate_index(index) + fn increment_core_instance_count(&self) -> Result<()> { + self.ondemand.increment_core_instance_count() } - fn allocate_memories( + fn decrement_core_instance_count(&self) { + self.ondemand.decrement_core_instance_count(); + } + + unsafe fn allocate_memory( &self, - index: usize, - req: &mut InstanceAllocationRequest, - mem: &mut PrimaryMap, - ) -> Result<()> { - assert_eq!(req.runtime_info.module().memory_plans.len(), 1); + request: &mut InstanceAllocationRequest, + memory_plan: &MemoryPlan, + memory_index: DefinedMemoryIndex, + ) -> Result<(MemoryAllocationIndex, Memory)> { + #[cfg(debug_assertions)] + { + let module = request.runtime_info.module(); + let offsets = request.runtime_info.offsets(); + self.validate_module_impl(module, offsets) + .expect("should have already validated the module before allocating memory"); + } + match self.preallocation { - Some(shared_memory) => { - mem.push(shared_memory.clone().as_memory()); - } - None => { - self.ondemand.allocate_memories(index, req, mem)?; - } + Some(shared_memory) => Ok(( + MemoryAllocationIndex::default(), + shared_memory.clone().as_memory(), + )), + None => self + .ondemand + .allocate_memory(request, memory_plan, memory_index), } - Ok(()) } - fn deallocate_memories(&self, index: usize, mems: &mut PrimaryMap) { - self.ondemand.deallocate_memories(index, mems) + unsafe fn deallocate_memory( + &self, + memory_index: DefinedMemoryIndex, + allocation_index: MemoryAllocationIndex, + memory: Memory, + ) { + self.ondemand + .deallocate_memory(memory_index, allocation_index, memory) } - fn allocate_tables( + unsafe fn allocate_table( &self, - index: usize, req: &mut InstanceAllocationRequest, - tables: &mut PrimaryMap, - ) -> Result<()> { - self.ondemand.allocate_tables(index, req, tables) + table_plan: &wasmtime_environ::TablePlan, + table_index: DefinedTableIndex, + ) -> Result<(TableAllocationIndex, Table)> { + self.ondemand.allocate_table(req, table_plan, table_index) } - fn deallocate_tables(&self, index: usize, tables: &mut PrimaryMap) { - self.ondemand.deallocate_tables(index, tables) + unsafe fn deallocate_table( + &self, + table_index: DefinedTableIndex, + allocation_index: TableAllocationIndex, + table: Table, + ) { + self.ondemand + .deallocate_table(table_index, allocation_index, table) } #[cfg(feature = "async")] diff --git a/fuzz/fuzz_targets/instantiate-many.rs b/fuzz/fuzz_targets/instantiate-many.rs index 9245071cad02..0a8fd9c92b6b 100644 --- a/fuzz/fuzz_targets/instantiate-many.rs +++ b/fuzz/fuzz_targets/instantiate-many.rs @@ -41,7 +41,7 @@ fn execute_one(data: &[u8]) -> Result<()> { let max_instances = match &config.wasmtime.strategy { generators::InstanceAllocationStrategy::OnDemand => u.int_in_range(1..=100)?, - generators::InstanceAllocationStrategy::Pooling(config) => config.instance_count, + generators::InstanceAllocationStrategy::Pooling(config) => config.total_core_instances, }; // Front-load with instantiation commands diff --git a/tests/all/async_functions.rs b/tests/all/async_functions.rs index 60ec31888e5e..f2cac9ffdd77 100644 --- a/tests/all/async_functions.rs +++ b/tests/all/async_functions.rs @@ -360,10 +360,8 @@ async fn fuel_eventually_finishes() { #[tokio::test] async fn async_with_pooling_stacks() { - let mut pool = PoolingAllocationConfig::default(); - pool.instance_count(1) - .instance_memory_pages(1) - .instance_table_elements(0); + let mut pool = crate::small_pool_config(); + pool.total_stacks(1).memory_pages(1).table_elements(0); let mut config = Config::new(); config.async_support(true); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); @@ -385,11 +383,8 @@ async fn async_with_pooling_stacks() { #[tokio::test] async fn async_host_func_with_pooling_stacks() -> Result<()> { - let mut pooling = PoolingAllocationConfig::default(); - pooling - .instance_count(1) - .instance_memory_pages(1) - .instance_table_elements(0); + let mut pooling = crate::small_pool_config(); + pooling.total_stacks(1).memory_pages(1).table_elements(0); let mut config = Config::new(); config.async_support(true); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pooling)); diff --git a/tests/all/instance.rs b/tests/all/instance.rs index 23d734a1cc39..abc7c8c29c0a 100644 --- a/tests/all/instance.rs +++ b/tests/all/instance.rs @@ -42,8 +42,8 @@ fn linear_memory_limits() -> Result<()> { return Ok(()); } test(&Engine::default())?; - let mut pool = PoolingAllocationConfig::default(); - pool.instance_memory_pages(65536); + let mut pool = crate::small_pool_config(); + pool.memory_pages(65536); test(&Engine::new(Config::new().allocation_strategy( InstanceAllocationStrategy::Pooling(pool), ))?)?; diff --git a/tests/all/limits.rs b/tests/all/limits.rs index 810902f11332..9adbb75877bc 100644 --- a/tests/all/limits.rs +++ b/tests/all/limits.rs @@ -352,8 +352,10 @@ fn test_initial_table_limits_exceeded() -> Result<()> { #[test] fn test_pooling_allocator_initial_limits_exceeded() -> Result<()> { - let mut pool = PoolingAllocationConfig::default(); - pool.instance_count(1).instance_memories(2); + let mut pool = crate::small_pool_config(); + pool.total_memories(2) + .max_memories_per_module(2) + .memory_pages(5); let mut config = Config::new(); config.wasm_multi_memory(true); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); @@ -740,8 +742,8 @@ fn custom_limiter_detect_grow_failure() -> Result<()> { if std::env::var("WASMTIME_TEST_NO_HOG_MEMORY").is_ok() { return Ok(()); } - let mut pool = PoolingAllocationConfig::default(); - pool.instance_memory_pages(10).instance_table_elements(10); + let mut pool = crate::small_pool_config(); + pool.memory_pages(10).table_elements(10); let mut config = Config::new(); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); let engine = Engine::new(&config).unwrap(); @@ -849,8 +851,8 @@ async fn custom_limiter_async_detect_grow_failure() -> Result<()> { if std::env::var("WASMTIME_TEST_NO_HOG_MEMORY").is_ok() { return Ok(()); } - let mut pool = PoolingAllocationConfig::default(); - pool.instance_memory_pages(10).instance_table_elements(10); + let mut pool = crate::small_pool_config(); + pool.memory_pages(10).table_elements(10); let mut config = Config::new(); config.async_support(true); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); diff --git a/tests/all/main.rs b/tests/all/main.rs index 671730b8d745..1c3e6d4af483 100644 --- a/tests/all/main.rs +++ b/tests/all/main.rs @@ -80,3 +80,19 @@ pub(crate) fn skip_pooling_allocator_tests() -> bool { // - https://github.com/bytecodealliance/wasmtime/pull/2518#issuecomment-747280133 std::env::var("WASMTIME_TEST_NO_HOG_MEMORY").is_ok() } + +/// Get the default pooling allocator configuration for tests, which is a +/// smaller pool than the normal default. +pub(crate) fn small_pool_config() -> wasmtime::PoolingAllocationConfig { + let mut config = wasmtime::PoolingAllocationConfig::default(); + + config.total_memories(1); + config.memory_pages(1); + config.total_tables(1); + config.table_elements(10); + + #[cfg(feature = "async")] + config.total_stacks(1); + + config +} diff --git a/tests/all/memory.rs b/tests/all/memory.rs index ecd1a0ab5881..3acdfc2c69c2 100644 --- a/tests/all/memory.rs +++ b/tests/all/memory.rs @@ -191,8 +191,8 @@ fn guards_present() -> Result<()> { fn guards_present_pooling() -> Result<()> { const GUARD_SIZE: u64 = 65536; - let mut pool = PoolingAllocationConfig::default(); - pool.instance_count(2).instance_memory_pages(10); + let mut pool = crate::small_pool_config(); + pool.total_memories(2).memory_pages(10); let mut config = Config::new(); config.static_memory_maximum_size(1 << 20); config.dynamic_memory_guard_size(GUARD_SIZE); diff --git a/tests/all/pooling_allocator.rs b/tests/all/pooling_allocator.rs index 004a2069b61d..64fc29b1ad26 100644 --- a/tests/all/pooling_allocator.rs +++ b/tests/all/pooling_allocator.rs @@ -4,10 +4,7 @@ use wasmtime::*; #[test] fn successful_instantiation() -> Result<()> { - let mut pool = PoolingAllocationConfig::default(); - pool.instance_count(1) - .instance_memory_pages(1) - .instance_table_elements(10); + let pool = crate::small_pool_config(); let mut config = Config::new(); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); config.dynamic_memory_guard_size(0); @@ -27,10 +24,8 @@ fn successful_instantiation() -> Result<()> { #[test] #[cfg_attr(miri, ignore)] fn memory_limit() -> Result<()> { - let mut pool = PoolingAllocationConfig::default(); - pool.instance_count(1) - .instance_memory_pages(3) - .instance_table_elements(10); + let mut pool = crate::small_pool_config(); + pool.memory_pages(3); let mut config = Config::new(); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); config.dynamic_memory_guard_size(0); @@ -45,7 +40,7 @@ fn memory_limit() -> Result<()> { Ok(_) => panic!("module instantiation should fail"), Err(e) => assert_eq!( e.to_string(), - "defined memories count of 2 exceeds the limit of 1", + "defined memories count of 2 exceeds the per-instance limit of 1", ), } @@ -102,10 +97,8 @@ fn memory_limit() -> Result<()> { #[test] fn memory_init() -> Result<()> { - let mut pool = PoolingAllocationConfig::default(); - pool.instance_count(1) - .instance_memory_pages(2) - .instance_table_elements(0); + let mut pool = crate::small_pool_config(); + pool.memory_pages(2).table_elements(0); let mut config = Config::new(); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); @@ -138,10 +131,8 @@ fn memory_init() -> Result<()> { #[test] #[cfg_attr(miri, ignore)] fn memory_guard_page_trap() -> Result<()> { - let mut pool = PoolingAllocationConfig::default(); - pool.instance_count(1) - .instance_memory_pages(2) - .instance_table_elements(0); + let mut pool = crate::small_pool_config(); + pool.memory_pages(2).table_elements(0); let mut config = Config::new(); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); @@ -207,10 +198,8 @@ fn memory_zeroed() -> Result<()> { return Ok(()); } - let mut pool = PoolingAllocationConfig::default(); - pool.instance_count(1) - .instance_memory_pages(1) - .instance_table_elements(0); + let mut pool = crate::small_pool_config(); + pool.memory_pages(1).table_elements(0); let mut config = Config::new(); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); config.dynamic_memory_guard_size(0); @@ -247,10 +236,8 @@ fn memory_zeroed() -> Result<()> { #[cfg_attr(miri, ignore)] fn table_limit() -> Result<()> { const TABLE_ELEMENTS: u32 = 10; - let mut pool = PoolingAllocationConfig::default(); - pool.instance_count(1) - .instance_memory_pages(1) - .instance_table_elements(TABLE_ELEMENTS); + let mut pool = crate::small_pool_config(); + pool.table_elements(TABLE_ELEMENTS); let mut config = Config::new(); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); config.dynamic_memory_guard_size(0); @@ -264,7 +251,7 @@ fn table_limit() -> Result<()> { Ok(_) => panic!("module compilation should fail"), Err(e) => assert_eq!( e.to_string(), - "defined tables count of 2 exceeds the limit of 1", + "defined tables count of 2 exceeds the per-instance limit of 1", ), } @@ -331,10 +318,8 @@ fn table_limit() -> Result<()> { #[test] #[cfg_attr(miri, ignore)] fn table_init() -> Result<()> { - let mut pool = PoolingAllocationConfig::default(); - pool.instance_count(1) - .instance_memory_pages(0) - .instance_table_elements(6); + let mut pool = crate::small_pool_config(); + pool.memory_pages(0).table_elements(6); let mut config = Config::new(); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); @@ -388,10 +373,7 @@ fn table_zeroed() -> Result<()> { return Ok(()); } - let mut pool = PoolingAllocationConfig::default(); - pool.instance_count(1) - .instance_memory_pages(1) - .instance_table_elements(10); + let pool = crate::small_pool_config(); let mut config = Config::new(); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); config.dynamic_memory_guard_size(0); @@ -426,12 +408,10 @@ fn table_zeroed() -> Result<()> { } #[test] -fn instantiation_limit() -> Result<()> { +fn total_core_instances_limit() -> Result<()> { const INSTANCE_LIMIT: u32 = 10; - let mut pool = PoolingAllocationConfig::default(); - pool.instance_count(INSTANCE_LIMIT) - .instance_memory_pages(1) - .instance_table_elements(10); + let mut pool = crate::small_pool_config(); + pool.total_core_instances(INSTANCE_LIMIT); let mut config = Config::new(); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); config.dynamic_memory_guard_size(0); @@ -454,7 +434,7 @@ fn instantiation_limit() -> Result<()> { Err(e) => assert_eq!( e.to_string(), format!( - "maximum concurrent instance limit of {} reached", + "maximum concurrent core instance limit of {} reached", INSTANCE_LIMIT ) ), @@ -474,10 +454,8 @@ fn instantiation_limit() -> Result<()> { #[test] fn preserve_data_segments() -> Result<()> { - let mut pool = PoolingAllocationConfig::default(); - pool.instance_count(2) - .instance_memory_pages(1) - .instance_table_elements(10); + let mut pool = crate::small_pool_config(); + pool.total_memories(2); let mut config = Config::new(); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); let engine = Engine::new(&config)?; @@ -524,10 +502,8 @@ fn multi_memory_with_imported_memories() -> Result<()> { // This test checks that the base address for the defined memory is correct for the instance // despite the presence of an imported memory. - let mut pool = PoolingAllocationConfig::default(); - pool.instance_count(1) - .instance_memories(2) - .instance_memory_pages(1); + let mut pool = crate::small_pool_config(); + pool.total_memories(2).max_memories_per_module(2); let mut config = Config::new(); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); config.wasm_multi_memory(true); @@ -565,8 +541,7 @@ fn drop_externref_global_during_module_init() -> Result<()> { } } - let mut pool = PoolingAllocationConfig::default(); - pool.instance_count(1); + let pool = crate::small_pool_config(); let mut config = Config::new(); config.wasm_reference_types(true); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); @@ -610,8 +585,7 @@ fn drop_externref_global_during_module_init() -> Result<()> { #[test] #[cfg_attr(miri, ignore)] fn switch_image_and_non_image() -> Result<()> { - let mut pool = PoolingAllocationConfig::default(); - pool.instance_count(1); + let pool = crate::small_pool_config(); let mut c = Config::new(); c.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); let engine = Engine::new(&c)?; @@ -669,8 +643,8 @@ fn switch_image_and_non_image() -> Result<()> { #[cfg(target_pointer_width = "64")] #[cfg_attr(miri, ignore)] fn instance_too_large() -> Result<()> { - let mut pool = PoolingAllocationConfig::default(); - pool.instance_size(16).instance_count(1); + let mut pool = crate::small_pool_config(); + pool.core_instance_size(16); let mut config = Config::new(); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); @@ -713,9 +687,8 @@ configured maximum of 16 bytes; breakdown of allocation requirement: fn dynamic_memory_pooling_allocator() -> Result<()> { for guard_size in [0, 1 << 16] { let max_size = 128 << 20; - let mut pool = PoolingAllocationConfig::default(); - pool.instance_count(1) - .instance_memory_pages(max_size / (64 * 1024)); + let mut pool = crate::small_pool_config(); + pool.memory_pages(max_size / (64 * 1024)); let mut config = Config::new(); config.static_memory_maximum_size(max_size); config.dynamic_memory_guard_size(guard_size); @@ -820,8 +793,8 @@ fn dynamic_memory_pooling_allocator() -> Result<()> { #[test] #[cfg_attr(miri, ignore)] fn zero_memory_pages_disallows_oob() -> Result<()> { - let mut pool = PoolingAllocationConfig::default(); - pool.instance_count(1).instance_memory_pages(0); + let mut pool = crate::small_pool_config(); + pool.memory_pages(0); let mut config = Config::new(); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); @@ -853,3 +826,360 @@ fn zero_memory_pages_disallows_oob() -> Result<()> { } Ok(()) } + +#[test] +#[cfg(feature = "component-model")] +fn total_component_instances_limit() -> Result<()> { + const TOTAL_COMPONENT_INSTANCES: u32 = 5; + + let mut pool = crate::small_pool_config(); + pool.total_component_instances(TOTAL_COMPONENT_INSTANCES); + let mut config = Config::new(); + config.wasm_component_model(true); + config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); + + let engine = Engine::new(&config)?; + let linker = wasmtime::component::Linker::new(&engine); + let component = wasmtime::component::Component::new(&engine, "(component)")?; + + let mut store = Store::new(&engine, ()); + for _ in 0..TOTAL_COMPONENT_INSTANCES { + linker.instantiate(&mut store, &component)?; + } + + match linker.instantiate(&mut store, &component) { + Ok(_) => panic!("should have hit component instance limit"), + Err(e) => assert_eq!( + e.to_string(), + format!( + "maximum concurrent component instance limit of {} reached", + TOTAL_COMPONENT_INSTANCES + ), + ), + } + + drop(store); + let mut store = Store::new(&engine, ()); + for _ in 0..TOTAL_COMPONENT_INSTANCES { + linker.instantiate(&mut store, &component)?; + } + + Ok(()) +} + +#[test] +#[cfg(feature = "component-model")] +fn component_instance_size_limit() -> Result<()> { + let mut pool = crate::small_pool_config(); + pool.component_instance_size(1); + let mut config = Config::new(); + config.wasm_component_model(true); + config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); + let engine = Engine::new(&config)?; + + match wasmtime::component::Component::new(&engine, "(component)") { + Ok(_) => panic!("should have hit limit"), + Err(e) => assert_eq!( + e.to_string(), + "instance allocation for this component requires 64 bytes of `VMComponentContext` space \ + which exceeds the configured maximum of 1 bytes" + ), + } + + Ok(()) +} + +#[test] +#[cfg_attr(miri, ignore)] +fn total_tables_limit() -> Result<()> { + const TOTAL_TABLES: u32 = 5; + + let mut pool = crate::small_pool_config(); + pool.total_tables(TOTAL_TABLES) + .total_core_instances(TOTAL_TABLES + 1); + let mut config = Config::new(); + config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); + + let engine = Engine::new(&config)?; + let linker = Linker::new(&engine); + let module = Module::new(&engine, "(module (table 0 1 funcref))")?; + + let mut store = Store::new(&engine, ()); + for _ in 0..TOTAL_TABLES { + linker.instantiate(&mut store, &module)?; + } + + match linker.instantiate(&mut store, &module) { + Ok(_) => panic!("should have hit table limit"), + Err(e) => assert_eq!( + e.to_string(), + format!("maximum concurrent table limit of {} reached", TOTAL_TABLES), + ), + } + + drop(store); + let mut store = Store::new(&engine, ()); + for _ in 0..TOTAL_TABLES { + linker.instantiate(&mut store, &module)?; + } + + Ok(()) +} + +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn total_stacks_limit() -> Result<()> { + use super::async_functions::PollOnce; + + const TOTAL_STACKS: u32 = 2; + + let mut pool = crate::small_pool_config(); + pool.total_stacks(TOTAL_STACKS) + .total_core_instances(TOTAL_STACKS + 1); + let mut config = Config::new(); + config.async_support(true); + config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); + + let engine = Engine::new(&config)?; + + let mut linker = Linker::new(&engine); + linker.func_new_async( + "async", + "yield", + FuncType::new([], []), + |_caller, _params, _results| { + Box::new(async { + tokio::task::yield_now().await; + Ok(()) + }) + }, + )?; + + let module = Module::new( + &engine, + r#" + (module + (import "async" "yield" (func $yield)) + (func (export "run") + call $yield + ) + ) + "#, + )?; + + // Allocate stacks up to the limit. (Poll the futures once to make sure we + // actually enter Wasm and force a stack allocation.) + + let mut store1 = Store::new(&engine, ()); + let instance1 = linker.instantiate_async(&mut store1, &module).await?; + let run1 = instance1.get_func(&mut store1, "run").unwrap(); + let future1 = PollOnce::new(Box::pin(run1.call_async(store1, &[], &mut []))) + .await + .unwrap_err(); + + let mut store2 = Store::new(&engine, ()); + let instance2 = linker.instantiate_async(&mut store2, &module).await?; + let run2 = instance2.get_func(&mut store2, "run").unwrap(); + let future2 = PollOnce::new(Box::pin(run2.call_async(store2, &[], &mut []))) + .await + .unwrap_err(); + + // Allocating more should fail. + let mut store3 = Store::new(&engine, ()); + match linker.instantiate_async(&mut store3, &module).await { + Ok(_) => panic!("should have hit stack limit"), + Err(e) => assert_eq!( + e.to_string(), + format!("maximum concurrent fiber limit of {} reached", TOTAL_STACKS), + ), + } + + // Finish the futures and return their Wasm stacks to the pool. + future1.await?; + future2.await?; + + // Should be able to allocate new stacks again. + let mut store1 = Store::new(&engine, ()); + let instance1 = linker.instantiate_async(&mut store1, &module).await?; + let run1 = instance1.get_func(&mut store1, "run").unwrap(); + let future1 = run1.call_async(store1, &[], &mut []); + + let mut store2 = Store::new(&engine, ()); + let instance2 = linker.instantiate_async(&mut store2, &module).await?; + let run2 = instance2.get_func(&mut store2, "run").unwrap(); + let future2 = run2.call_async(store2, &[], &mut []); + + future1.await?; + future2.await?; + + Ok(()) +} + +#[test] +#[cfg(feature = "component-model")] +fn component_core_instances_limit() -> Result<()> { + let mut pool = crate::small_pool_config(); + pool.max_core_instances_per_component(1); + let mut config = Config::new(); + config.wasm_component_model(true); + config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); + let engine = Engine::new(&config)?; + + // One core instance works. + wasmtime::component::Component::new( + &engine, + r#" + (component + (core module $m) + (core instance $a (instantiate $m)) + ) + "#, + )?; + + // Two core instances doesn't. + match wasmtime::component::Component::new( + &engine, + r#" + (component + (core module $m) + (core instance $a (instantiate $m)) + (core instance $b (instantiate $m)) + ) + "#, + ) { + Ok(_) => panic!("should have hit limit"), + Err(e) => assert_eq!( + e.to_string(), + "The component transitively contains 2 core module instances, which exceeds the \ + configured maximum of 1" + ), + } + + Ok(()) +} + +#[test] +#[cfg(feature = "component-model")] +fn component_memories_limit() -> Result<()> { + let mut pool = crate::small_pool_config(); + pool.max_memories_per_component(1).total_memories(2); + let mut config = Config::new(); + config.wasm_component_model(true); + config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); + let engine = Engine::new(&config)?; + + // One memory works. + wasmtime::component::Component::new( + &engine, + r#" + (component + (core module $m (memory 1 1)) + (core instance $a (instantiate $m)) + ) + "#, + )?; + + // Two memories doesn't. + match wasmtime::component::Component::new( + &engine, + r#" + (component + (core module $m (memory 1 1)) + (core instance $a (instantiate $m)) + (core instance $b (instantiate $m)) + ) + "#, + ) { + Ok(_) => panic!("should have hit limit"), + Err(e) => assert_eq!( + e.to_string(), + "The component transitively contains 2 Wasm linear memories, which exceeds the \ + configured maximum of 1" + ), + } + + Ok(()) +} + +#[test] +#[cfg(feature = "component-model")] +fn component_tables_limit() -> Result<()> { + let mut pool = crate::small_pool_config(); + pool.max_tables_per_component(1).total_tables(2); + let mut config = Config::new(); + config.wasm_component_model(true); + config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); + let engine = Engine::new(&config)?; + + // One table works. + wasmtime::component::Component::new( + &engine, + r#" + (component + (core module $m (table 1 1 funcref)) + (core instance $a (instantiate $m)) + ) + "#, + )?; + + // Two tables doesn't. + match wasmtime::component::Component::new( + &engine, + r#" + (component + (core module $m (table 1 1 funcref)) + (core instance $a (instantiate $m)) + (core instance $b (instantiate $m)) + ) + "#, + ) { + Ok(_) => panic!("should have hit limit"), + Err(e) => assert_eq!( + e.to_string(), + "The component transitively contains 2 tables, which exceeds the \ + configured maximum of 1" + ), + } + + Ok(()) +} + +#[test] +#[cfg_attr(miri, ignore)] +fn total_memories_limit() -> Result<()> { + const TOTAL_MEMORIES: u32 = 5; + + let mut pool = crate::small_pool_config(); + pool.total_memories(TOTAL_MEMORIES) + .total_core_instances(TOTAL_MEMORIES + 1); + let mut config = Config::new(); + config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); + + let engine = Engine::new(&config)?; + let linker = Linker::new(&engine); + let module = Module::new(&engine, "(module (memory 1 1))")?; + + let mut store = Store::new(&engine, ()); + for _ in 0..TOTAL_MEMORIES { + linker.instantiate(&mut store, &module)?; + } + + match linker.instantiate(&mut store, &module) { + Ok(_) => panic!("should have hit memory limit"), + Err(e) => assert_eq!( + e.to_string(), + format!( + "maximum concurrent memory limit of {} reached", + TOTAL_MEMORIES + ), + ), + } + + drop(store); + let mut store = Store::new(&engine, ()); + for _ in 0..TOTAL_MEMORIES { + linker.instantiate(&mut store, &module)?; + } + + Ok(()) +} diff --git a/tests/all/wast.rs b/tests/all/wast.rs index 1dbd4371d7cf..8f2f357be67d 100644 --- a/tests/all/wast.rs +++ b/tests/all/wast.rs @@ -127,10 +127,10 @@ fn run_wast(wast: &str, strategy: Strategy, pooling: bool) -> anyhow::Result<()> // If a wast test fails because of a limit being "exceeded" or if memory/table // fails to grow, the values here will need to be adjusted. let mut pool = PoolingAllocationConfig::default(); - pool.instance_count(450) - .instance_memories(if multi_memory { 9 } else { 1 }) - .instance_tables(4) - .instance_memory_pages(805); + pool.total_memories(450) + .memory_pages(805) + .max_memories_per_module(if multi_memory { 9 } else { 1 }) + .max_tables_per_module(4); cfg.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); Some(lock_pooling()) } else { From bb64a534dad04594716ae81eb2dcd1da07a93726 Mon Sep 17 00:00:00 2001 From: Nick Fitzgerald Date: Wed, 16 Aug 2023 11:14:19 -0700 Subject: [PATCH 04/11] Remove unused file Messed up from rebasing, this code is actually just inline in the index allocator module. --- .../src/instance/allocator/pooling/simple.rs | 24 ------------------- 1 file changed, 24 deletions(-) delete mode 100644 crates/runtime/src/instance/allocator/pooling/simple.rs diff --git a/crates/runtime/src/instance/allocator/pooling/simple.rs b/crates/runtime/src/instance/allocator/pooling/simple.rs deleted file mode 100644 index f818d749b69c..000000000000 --- a/crates/runtime/src/instance/allocator/pooling/simple.rs +++ /dev/null @@ -1,24 +0,0 @@ -//! A simple index allocator. -//! -//! This index allocator doesn't do any module affinity or anything like that, -//! however it is built on top of the `ModuleAffinityIndexAllocator` to save -//! code (and code size). - -use super::module_affinity::{ModuleAffinityIndexAllocator, SlotId}; - -#[derive(Debug)] -pub struct SimpleIndexAllocator(ModuleAffinityIndexAllocator); - -impl SimpleIndexAllocator { - pub fn new(max_instances: u32) -> Self { - SimpleIndexAllocator(ModuleAffinityIndexAllocator::new(max_instances, 0)) - } - - pub fn alloc(&self) -> Option { - self.0.alloc(None) - } - - pub(crate) fn free(&self, index: SlotId) { - self.0.free(index); - } -} From c4aa70ee74bc2350b9e3a41ec8810f6a4e422673 Mon Sep 17 00:00:00 2001 From: Nick Fitzgerald Date: Wed, 16 Aug 2023 11:50:17 -0700 Subject: [PATCH 05/11] Address review feedback --- .../fuzzing/src/generators/pooling_config.rs | 4 +- crates/runtime/src/instance/allocator.rs | 43 ++++++++---- .../allocator/pooling/index_allocator.rs | 68 ++++++++++--------- .../instance/allocator/pooling/memory_pool.rs | 28 +++++--- crates/wasmtime/src/component/component.rs | 16 ++--- crates/wasmtime/src/config.rs | 30 ++++---- crates/wasmtime/src/store.rs | 13 ++-- tests/all/pooling_allocator.rs | 4 +- 8 files changed, 120 insertions(+), 86 deletions(-) diff --git a/crates/fuzzing/src/generators/pooling_config.rs b/crates/fuzzing/src/generators/pooling_config.rs index 4eead14cad56..9cdbd1d1319d 100644 --- a/crates/fuzzing/src/generators/pooling_config.rs +++ b/crates/fuzzing/src/generators/pooling_config.rs @@ -46,11 +46,11 @@ impl PoolingAllocationConfig { cfg.memory_pages(self.memory_pages); cfg.table_elements(self.table_elements); - cfg.component_instance_size(self.component_instance_size); + cfg.max_component_instance_size(self.component_instance_size); cfg.max_memories_per_component(self.max_memories_per_component); cfg.max_tables_per_component(self.max_tables_per_component); - cfg.core_instance_size(self.core_instance_size); + cfg.max_core_instance_size(self.core_instance_size); cfg.max_memories_per_module(self.max_memories_per_module); cfg.max_tables_per_module(self.max_tables_per_module); diff --git a/crates/runtime/src/instance/allocator.rs b/crates/runtime/src/instance/allocator.rs index 5388efb4ed2e..6a17e14682f9 100644 --- a/crates/runtime/src/instance/allocator.rs +++ b/crates/runtime/src/instance/allocator.rs @@ -163,6 +163,19 @@ pub unsafe trait InstanceAllocatorImpl { /// Not all instance allocators will have limits for the maximum number of /// concurrent component instances that can be live at the same time, and /// these allocators may implement this method with a no-op. + // + // Note: It would be nice to have an associated type that on construction + // does the increment and on drop does the decrement but there are two + // problems with this: + // + // 1. This trait's implementations are always used as trait objects, and + // associated types are not object safe. + // + // 2. We would want a parameterized `Drop` implementation so that we could + // pass in the `InstaceAllocatorImpl` on drop, but this doesn't exist in + // Rust. Therefore, we would be forced to add reference counting and + // stuff like that to keep a handle on the instance allocator from this + // theoretical type. That's a bummer. fn increment_component_instance_count(&self) -> Result<()>; /// The dual of `increment_component_instance_count`. @@ -308,23 +321,23 @@ pub trait InstanceAllocator: InstanceAllocatorImpl { let num_defined_tables = module.table_plans.len() - module.num_imported_tables; let mut tables = PrimaryMap::with_capacity(num_defined_tables); - let result = self - .allocate_memories(&mut request, &mut memories) - .and_then(|()| self.allocate_tables(&mut request, &mut tables)); - if let Err(e) = result { - self.deallocate_memories(&mut memories); - self.deallocate_tables(&mut tables); - self.decrement_core_instance_count(); - return Err(e); - } - - unsafe { - Ok(Instance::new( + match (|| { + self.allocate_memories(&mut request, &mut memories)?; + self.allocate_tables(&mut request, &mut tables)?; + Ok(()) + })() { + Ok(_) => Ok(Instance::new( request, memories, tables, &module.memory_plans, - )) + )), + Err(e) => { + self.deallocate_memories(&mut memories); + self.deallocate_tables(&mut tables); + self.decrement_core_instance_count(); + Err(e) + } } } @@ -392,6 +405,10 @@ pub trait InstanceAllocator: InstanceAllocatorImpl { memories: &mut PrimaryMap, ) { for (memory_index, (allocation_index, memory)) in mem::take(memories) { + // Because deallocating memory is infallible, we don't need to worry + // about leaking subsequent memories if the first memory failed to + // deallocate. If deallocating memory ever becomes fallible, we will + // need to be careful here! self.deallocate_memory(memory_index, allocation_index, memory); } } diff --git a/crates/runtime/src/instance/allocator/pooling/index_allocator.rs b/crates/runtime/src/instance/allocator/pooling/index_allocator.rs index ec4dbf1060f8..5ae2f862a6c9 100644 --- a/crates/runtime/src/instance/allocator/pooling/index_allocator.rs +++ b/crates/runtime/src/instance/allocator/pooling/index_allocator.rs @@ -44,6 +44,10 @@ impl SimpleIndexAllocator { } } +/// A particular defined memory within a particular module. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub struct MemoryInModule(pub CompiledModuleId, pub DefinedMemoryIndex); + /// An index allocator that has configurable affinity between slots and modules /// so that slots are often reused for the same module again. #[derive(Debug)] @@ -86,7 +90,7 @@ struct Inner { /// /// The `List` here is appended to during deallocation and removal happens /// from the tail during allocation. - module_affine: HashMap<(CompiledModuleId, DefinedMemoryIndex), List>, + module_affine: HashMap, } /// A helper "linked list" data structure which is based on indices. @@ -107,7 +111,7 @@ struct Link { #[derive(Clone, Debug)] enum SlotState { /// This slot is currently in use and is affine to the specified module's memory. - Used(Option<(CompiledModuleId, DefinedMemoryIndex)>), + Used(Option), /// This slot is not currently used, and has never been used. UnusedCold, @@ -131,7 +135,7 @@ impl SlotState { #[derive(Default, Copy, Clone, Debug)] struct Unused { /// Which module this slot was historically affine to, if any. - affinity: Option<(CompiledModuleId, DefinedMemoryIndex)>, + affinity: Option, /// Metadata about the linked list for all slots affine to `affinity`. affine_list_link: Link, @@ -162,10 +166,7 @@ impl ModuleAffinityIndexAllocator { /// affinity request if the allocation strategy supports it. /// /// Returns `None` if no more slots are available. - pub fn alloc( - &self, - for_memory: Option<(CompiledModuleId, DefinedMemoryIndex)>, - ) -> Option { + pub fn alloc(&self, for_memory: Option) -> Option { self._alloc(for_memory, AllocMode::AnySlot) } @@ -182,16 +183,12 @@ impl ModuleAffinityIndexAllocator { memory_index: DefinedMemoryIndex, ) -> Option { self._alloc( - Some((module_id, memory_index)), + Some(MemoryInModule(module_id, memory_index)), AllocMode::ForceAffineAndClear, ) } - fn _alloc( - &self, - for_memory: Option<(CompiledModuleId, DefinedMemoryIndex)>, - mode: AllocMode, - ) -> Option { + fn _alloc(&self, for_memory: Option, mode: AllocMode) -> Option { let mut inner = self.0.lock().unwrap(); let inner = &mut *inner; @@ -302,9 +299,7 @@ impl ModuleAffinityIndexAllocator { /// For testing only, get the list of all modules with at least /// one slot with affinity for that module. #[cfg(test)] - pub(crate) fn testing_module_affinity_list( - &self, - ) -> Vec<(CompiledModuleId, DefinedMemoryIndex)> { + pub(crate) fn testing_module_affinity_list(&self) -> Vec { let inner = self.0.lock().unwrap(); inner.module_affine.keys().copied().collect() } @@ -313,10 +308,7 @@ impl ModuleAffinityIndexAllocator { impl Inner { /// Attempts to allocate a slot already affine to `id`, returning `None` if /// `id` is `None` or if there are no affine slots. - fn pick_affine( - &mut self, - for_memory: Option<(CompiledModuleId, DefinedMemoryIndex)>, - ) -> Option { + fn pick_affine(&mut self, for_memory: Option) -> Option { // Note that the `tail` is chosen here of the affine list as it's the // most recently used, which for affine allocations is what we want -- // maximizing temporal reuse. @@ -477,8 +469,8 @@ mod test { #[test] fn test_affinity_allocation_strategy() { let id_alloc = CompiledModuleIdAllocator::new(); - let id1 = (id_alloc.alloc(), DefinedMemoryIndex::new(0)); - let id2 = (id_alloc.alloc(), DefinedMemoryIndex::new(0)); + let id1 = MemoryInModule(id_alloc.alloc(), DefinedMemoryIndex::new(0)); + let id2 = MemoryInModule(id_alloc.alloc(), DefinedMemoryIndex::new(0)); let state = ModuleAffinityIndexAllocator::new(100, 100); let index1 = state.alloc(Some(id1)).unwrap(); @@ -536,8 +528,8 @@ mod test { for max_unused_warm_slots in [0, 1, 2] { let state = ModuleAffinityIndexAllocator::new(100, max_unused_warm_slots); - let index1 = state.alloc(Some((id, memory_index))).unwrap(); - let index2 = state.alloc(Some((id, memory_index))).unwrap(); + let index1 = state.alloc(Some(MemoryInModule(id, memory_index))).unwrap(); + let index2 = state.alloc(Some(MemoryInModule(id, memory_index))).unwrap(); state.free(index2); state.free(index1); assert!(state @@ -559,9 +551,10 @@ mod test { let mut rng = rand::thread_rng(); let id_alloc = CompiledModuleIdAllocator::new(); - let ids = std::iter::repeat_with(|| (id_alloc.alloc(), DefinedMemoryIndex::new(0))) - .take(10) - .collect::>(); + let ids = + std::iter::repeat_with(|| MemoryInModule(id_alloc.alloc(), DefinedMemoryIndex::new(0))) + .take(10) + .collect::>(); let state = ModuleAffinityIndexAllocator::new(1000, 1000); let mut allocated: Vec = vec![]; let mut last_id = vec![None; 1000]; @@ -603,9 +596,9 @@ mod test { #[test] fn test_affinity_threshold() { let id_alloc = CompiledModuleIdAllocator::new(); - let id1 = (id_alloc.alloc(), DefinedMemoryIndex::new(0)); - let id2 = (id_alloc.alloc(), DefinedMemoryIndex::new(0)); - let id3 = (id_alloc.alloc(), DefinedMemoryIndex::new(0)); + let id1 = MemoryInModule(id_alloc.alloc(), DefinedMemoryIndex::new(0)); + let id2 = MemoryInModule(id_alloc.alloc(), DefinedMemoryIndex::new(0)); + let id3 = MemoryInModule(id_alloc.alloc(), DefinedMemoryIndex::new(0)); let state = ModuleAffinityIndexAllocator::new(10, 2); // Set some slot affinities @@ -640,19 +633,28 @@ mod test { // LRU is 1, so that should be picked assert_eq!( - state.alloc(Some((id_alloc.alloc(), DefinedMemoryIndex::new(0)))), + state.alloc(Some(MemoryInModule( + id_alloc.alloc(), + DefinedMemoryIndex::new(0) + ))), Some(SlotId(1)) ); // Pick another LRU entry, this time 2 assert_eq!( - state.alloc(Some((id_alloc.alloc(), DefinedMemoryIndex::new(0)))), + state.alloc(Some(MemoryInModule( + id_alloc.alloc(), + DefinedMemoryIndex::new(0) + ))), Some(SlotId(2)) ); // This should preserve slot `0` and pick up something new assert_eq!( - state.alloc(Some((id_alloc.alloc(), DefinedMemoryIndex::new(0)))), + state.alloc(Some(MemoryInModule( + id_alloc.alloc(), + DefinedMemoryIndex::new(0) + ))), Some(SlotId(3)) ); diff --git a/crates/runtime/src/instance/allocator/pooling/memory_pool.rs b/crates/runtime/src/instance/allocator/pooling/memory_pool.rs index ba72e05b90c2..e215d59d5457 100644 --- a/crates/runtime/src/instance/allocator/pooling/memory_pool.rs +++ b/crates/runtime/src/instance/allocator/pooling/memory_pool.rs @@ -1,5 +1,5 @@ use super::{ - index_allocator::{ModuleAffinityIndexAllocator, SlotId}, + index_allocator::{MemoryInModule, ModuleAffinityIndexAllocator, SlotId}, MemoryAllocationIndex, }; use crate::{ @@ -17,9 +17,6 @@ use wasmtime_environ::{ /// /// A linear memory is divided into accessible pages and guard pages. /// -/// Each instance index into the pool returns an iterator over the base -/// addresses of the instance's linear memories. -/// /// A diagram for this struct's fields is: /// /// ```ignore @@ -35,7 +32,7 @@ use wasmtime_environ::{ /// +-----------+--------+---+-----------+ +--------+---+-----------+ /// | |<------------------+----------------------------------> /// \ | \ -/// mapping | `max_instances * max_memories` memories +/// mapping | `max_total_memories` memories /// / /// initial_memory_offset /// ``` @@ -65,6 +62,8 @@ pub struct MemoryPool { max_total_memories: usize, // The maximum number of memories that a single core module instance may // use. + // + // NB: this is needed for validation but does not affect the pool's size. memories_per_instance: usize, // How much linear memory, in bytes, to keep resident after resetting for // use with the next instance. This much memory will be `memset` to zero @@ -114,9 +113,8 @@ impl MemoryPool { 0 }; - // The entire allocation here is the size of each memory times the - // max memories per instance times the number of instances allowed in - // this pool, plus guard regions. + // The entire allocation here is the size of each memory (with guard + // regions) times the total number of memories in the pool. // // Note, though, that guard regions are required to be after each linear // memory. If the `guard_before_linear_memory` setting is specified, @@ -214,7 +212,7 @@ impl MemoryPool { request .runtime_info .unique_id() - .map(|id| (id, memory_index)), + .map(|id| MemoryInModule(id, memory_index)), ) .map(|slot| MemoryAllocationIndex(u32::try_from(slot.index()).unwrap())) .ok_or_else(|| { @@ -302,6 +300,18 @@ impl MemoryPool { // allocated further (the module is being dropped) so this shouldn't hit // any sort of infinite loop since this should be the final operation // working with `module`. + // + // TODO: We are given a module id, but key affinity by pair of module id + // and defined memory index. We are missing any defined memory index or + // count of how many memories the module defines here. Therefore, we + // probe up to the maximum number of memories per instance. This is fine + // because that maximum is generally relatively small. If this method + // somehow ever gets hot because of unnecessary probing, we should + // either pass in the actual number of defined memories for the given + // module to this method, or keep a side table of all slots that are + // associated with a module (not just module and memory). The latter + // would require care to make sure that its maintenance wouldn't be too + // expensive for normal allocation/free operations. for i in 0..self.memories_per_instance { use wasmtime_environ::EntityRef; let memory_index = DefinedMemoryIndex::new(i); diff --git a/crates/wasmtime/src/component/component.rs b/crates/wasmtime/src/component/component.rs index 6e4034a12e20..1b865f7c50db 100644 --- a/crates/wasmtime/src/component/component.rs +++ b/crates/wasmtime/src/component/component.rs @@ -247,6 +247,14 @@ impl Component { None => bincode::deserialize(code_memory.wasmtime_info())?, }; + // Validate that the component can be used with the current instance + // allocator. + engine.allocator().validate_component( + &info.component, + &VMComponentOffsets::new(HostPtr, &info.component), + &|module_index| &static_modules[module_index].module, + )?; + // Create a signature registration with the `Engine` for all trampolines // and core wasm types found within this component, both for the // component and for all included core wasm modules. @@ -258,14 +266,6 @@ impl Component { let types = Arc::new(types); let code = Arc::new(CodeObject::new(code_memory, signatures, types.into())); - // Validate that the component can be used with the current instance - // allocator. - engine.allocator().validate_component( - &info.component, - &VMComponentOffsets::new(HostPtr, &info.component), - &|module_index| &static_modules[module_index].module, - )?; - // Convert all information about static core wasm modules into actual // `Module` instances by converting each `CompiledModuleInfo`, the // `types` type information, and the code memory to a runtime object. diff --git a/crates/wasmtime/src/config.rs b/crates/wasmtime/src/config.rs index 9439baba3e84..d52137e94251 100644 --- a/crates/wasmtime/src/config.rs +++ b/crates/wasmtime/src/config.rs @@ -2027,13 +2027,13 @@ impl PoolingAllocationConfig { /// /// This provides an upper-bound on the total size of component /// metadata-related allocations, along with - /// [`PoolingAllocationConfig::component_instance_size`]. The upper bound is + /// [`PoolingAllocationConfig::max_component_instance_size`]. The upper bound is /// /// ```text - /// total_component_instances * component_instance_size + /// total_component_instances * max_component_instance_size /// ``` /// - /// where `component_instance_size` is rounded up to the size and alignment + /// where `max_component_instance_size` is rounded up to the size and alignment /// of the internal representation of the metadata. pub fn total_component_instances(&mut self, count: u32) -> &mut Self { self.config.limits.total_component_instances = count; @@ -2065,12 +2065,12 @@ impl PoolingAllocationConfig { /// [`PoolingAllocationConfig::total_component_instances`]. The upper bound is /// /// ```text - /// total_component_instances * component_instance_size + /// total_component_instances * max_component_instance_size /// ``` /// - /// where `component_instance_size` is rounded up to the size and alignment + /// where `max_component_instance_size` is rounded up to the size and alignment /// of the internal representation of the metadata. - pub fn component_instance_size(&mut self, size: usize) -> &mut Self { + pub fn max_component_instance_size(&mut self, size: usize) -> &mut Self { self.config.limits.component_instance_size = size; self } @@ -2081,7 +2081,7 @@ impl PoolingAllocationConfig { /// This method (along with /// [`PoolingAllocationConfig::max_memories_per_component`], /// [`PoolingAllocationConfig::max_tables_per_component`], and - /// [`PoolingAllocationConfig::component_instance_size`]) allows you to cap + /// [`PoolingAllocationConfig::max_component_instance_size`]) allows you to cap /// the amount of resources a single component allocation consumes. /// /// If a component will instantiate more core instances than `count`, then @@ -2097,7 +2097,7 @@ impl PoolingAllocationConfig { /// This method (along with /// [`PoolingAllocationConfig::max_core_instances_per_component`], /// [`PoolingAllocationConfig::max_tables_per_component`], and - /// [`PoolingAllocationConfig::component_instance_size`]) allows you to cap + /// [`PoolingAllocationConfig::max_component_instance_size`]) allows you to cap /// the amount of resources a single component allocation consumes. /// /// If a component transitively contains more linear memories than `count`, @@ -2113,7 +2113,7 @@ impl PoolingAllocationConfig { /// This method (along with /// [`PoolingAllocationConfig::max_core_instances_per_component`], /// [`PoolingAllocationConfig::max_memories_per_component`], - /// [`PoolingAllocationConfig::component_instance_size`]) allows you to cap + /// [`PoolingAllocationConfig::max_component_instance_size`]) allows you to cap /// the amount of resources a single component allocation consumes. /// /// If a component will transitively contains more tables than `count`, then @@ -2176,13 +2176,13 @@ impl PoolingAllocationConfig { /// /// This provides an upper-bound on the total size of core instance /// metadata-related allocations, along with - /// [`PoolingAllocationConfig::core_instance_size`]. The upper bound is + /// [`PoolingAllocationConfig::max_core_instance_size`]. The upper bound is /// /// ```text - /// total_core_instances * core_instance_size + /// total_core_instances * max_core_instance_size /// ``` /// - /// where `core_instance_size` is rounded up to the size and alignment of + /// where `max_core_instance_size` is rounded up to the size and alignment of /// the internal representation of the metadata. pub fn total_core_instances(&mut self, count: u32) -> &mut Self { self.config.limits.total_core_instances = count; @@ -2214,12 +2214,12 @@ impl PoolingAllocationConfig { /// [`PoolingAllocationConfig::total_core_instances`]. The upper bound is /// /// ```text - /// total_core_instances * core_instance_size + /// total_core_instances * max_core_instance_size /// ``` /// - /// where `core_instance_size` is rounded up to the size and alignment of + /// where `max_core_instance_size` is rounded up to the size and alignment of /// the internal representation of the metadata. - pub fn core_instance_size(&mut self, size: usize) -> &mut Self { + pub fn max_core_instance_size(&mut self, size: usize) -> &mut Self { self.config.limits.core_instance_size = size; self } diff --git a/crates/wasmtime/src/store.rs b/crates/wasmtime/src/store.rs index 1df0e8c7ab28..9da3b1e62428 100644 --- a/crates/wasmtime/src/store.rs +++ b/crates/wasmtime/src/store.rs @@ -292,7 +292,7 @@ pub struct StoreOpaque { runtime_limits: VMRuntimeLimits, instances: Vec, #[cfg(feature = "component-model")] - component_instances: Vec, + num_component_instances: usize, signal_handler: Option>>, externref_activations_table: VMExternRefActivationsTable, modules: ModuleRegistry, @@ -464,7 +464,7 @@ impl Store { runtime_limits: Default::default(), instances: Vec::new(), #[cfg(feature = "component-model")] - component_instances: Vec::new(), + num_component_instances: 0, signal_handler: None, externref_activations_table: VMExternRefActivationsTable::new(), modules: ModuleRegistry::default(), @@ -1592,7 +1592,12 @@ at https://bytecodealliance.org/security. #[cfg(feature = "component-model")] pub(crate) fn push_component_instance(&mut self, instance: crate::component::Instance) { - self.component_instances.push(instance); + // We don't actually need the instance itself right now, but it seems + // like something we will almost certainly eventually want to keep + // around, so force callers to provide it. + let _ = instance; + + self.num_component_instances += 1; } } @@ -2205,7 +2210,7 @@ impl Drop for StoreOpaque { #[cfg(feature = "component-model")] { - for _ in self.component_instances.iter() { + for _ in 0..self.num_component_instances { allocator.decrement_component_instance_count(); } } diff --git a/tests/all/pooling_allocator.rs b/tests/all/pooling_allocator.rs index 64fc29b1ad26..cd89d73267bf 100644 --- a/tests/all/pooling_allocator.rs +++ b/tests/all/pooling_allocator.rs @@ -644,7 +644,7 @@ fn switch_image_and_non_image() -> Result<()> { #[cfg_attr(miri, ignore)] fn instance_too_large() -> Result<()> { let mut pool = crate::small_pool_config(); - pool.core_instance_size(16); + pool.max_core_instance_size(16); let mut config = Config::new(); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); @@ -871,7 +871,7 @@ fn total_component_instances_limit() -> Result<()> { #[cfg(feature = "component-model")] fn component_instance_size_limit() -> Result<()> { let mut pool = crate::small_pool_config(); - pool.component_instance_size(1); + pool.max_component_instance_size(1); let mut config = Config::new(); config.wasm_component_model(true); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); From f14b6ceb1700b207c8cfa174d2cbdeb094bd9ba1 Mon Sep 17 00:00:00 2001 From: Nick Fitzgerald Date: Wed, 16 Aug 2023 14:13:14 -0700 Subject: [PATCH 06/11] Fix benchmarks build --- benches/thread_eager_init.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/benches/thread_eager_init.rs b/benches/thread_eager_init.rs index 8572a335c406..421ac043b5cb 100644 --- a/benches/thread_eager_init.rs +++ b/benches/thread_eager_init.rs @@ -92,7 +92,9 @@ fn test_setup() -> (Engine, Module) { let pool_count = 10; let mut pool = PoolingAllocationConfig::default(); - pool.instance_count(pool_count).instance_memory_pages(1); + pool.total_memories(pool_count) + .total_stacks(pool_count) + .total_tables(pool_count); let mut config = Config::new(); config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); let engine = Engine::new(&config).unwrap(); From 98fcc39630b78a1d6121589c8d99a3e966b326ff Mon Sep 17 00:00:00 2001 From: Nick Fitzgerald Date: Thu, 17 Aug 2023 08:59:00 -0700 Subject: [PATCH 07/11] Fix ignoring test under miri The `async_functions` module is not even compiled-but-ignored with miri, it is completely `cfg`ed off. Therefore we ahve to do the same with this test that imports stuff from that module. --- tests/all/pooling_allocator.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/all/pooling_allocator.rs b/tests/all/pooling_allocator.rs index cd89d73267bf..f6fcc2a541c0 100644 --- a/tests/all/pooling_allocator.rs +++ b/tests/all/pooling_allocator.rs @@ -927,7 +927,7 @@ fn total_tables_limit() -> Result<()> { } #[tokio::test] -#[cfg_attr(miri, ignore)] +#[cfg(not(miri))] async fn total_stacks_limit() -> Result<()> { use super::async_functions::PollOnce; From 393b18b030b7775ea552bf338816e3575efc729a Mon Sep 17 00:00:00 2001 From: Nick Fitzgerald Date: Thu, 17 Aug 2023 10:59:41 -0700 Subject: [PATCH 08/11] Fix doc links --- crates/wasmtime/src/config.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/crates/wasmtime/src/config.rs b/crates/wasmtime/src/config.rs index d52137e94251..dba6281ebee4 100644 --- a/crates/wasmtime/src/config.rs +++ b/crates/wasmtime/src/config.rs @@ -1929,7 +1929,7 @@ impl PoolingAllocationConfig { /// If this setting is set to infinity, however, then cold slots are /// prioritized to be allocated from. This means that the set of slots used /// over the lifetime of a program will approach - /// [`PoolingAllocationConfig::instance_count`], or the maximum number of + /// [`PoolingAllocationConfig::total_memories`], or the maximum number of /// slots in the pooling allocator. /// /// Wasmtime does not aggressively decommit all resources associated with a @@ -2192,10 +2192,10 @@ impl PoolingAllocationConfig { /// The maximum size, in bytes, allocated for a core instance's `VMContext` /// metadata. /// - /// The [`Instance`] type has a static size but its `VMContext` metadata is - /// dynamically sized depending on the module being instantiated. This size - /// limit loosely correlates to the size of the Wasm module, taking into - /// account factors such as: + /// The [`Instance`][crate::Instance] type has a static size but its + /// `VMContext` metadata is dynamically sized depending on the module being + /// instantiated. This size limit loosely correlates to the size of the Wasm + /// module, taking into account factors such as: /// /// * number of functions /// * number of globals From 0ef29635ffb86f92b55600ca5f40c70b6580e8cf Mon Sep 17 00:00:00 2001 From: Nick Fitzgerald Date: Thu, 17 Aug 2023 12:18:21 -0700 Subject: [PATCH 09/11] Allow testing utilities to be unused The exact `cfg`s that unlock the tests that use these are platform and feature dependent and ends up being like 5 things and super long. Simpler to just allow unused for when we are testing on other platforms or don't have the compile time features enabled. --- .../runtime/src/instance/allocator/pooling/index_allocator.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/crates/runtime/src/instance/allocator/pooling/index_allocator.rs b/crates/runtime/src/instance/allocator/pooling/index_allocator.rs index 5ae2f862a6c9..c7acf5d97d6b 100644 --- a/crates/runtime/src/instance/allocator/pooling/index_allocator.rs +++ b/crates/runtime/src/instance/allocator/pooling/index_allocator.rs @@ -39,6 +39,7 @@ impl SimpleIndexAllocator { } #[cfg(test)] + #[allow(unused)] pub(crate) fn testing_freelist(&self) -> Vec { self.0.testing_freelist() } @@ -288,6 +289,7 @@ impl ModuleAffinityIndexAllocator { /// For testing only, we want to be able to assert what is on the /// single freelist, for the policies that keep just one. #[cfg(test)] + #[allow(unused)] pub(crate) fn testing_freelist(&self) -> Vec { let inner = self.0.lock().unwrap(); inner @@ -424,6 +426,7 @@ impl List { } #[cfg(test)] + #[allow(unused)] fn iter<'a>( &'a self, states: &'a [SlotState], From c20ae756ad25c74a9ab973347b9e48bd7cda0b08 Mon Sep 17 00:00:00 2001 From: Nick Fitzgerald Date: Fri, 18 Aug 2023 12:50:40 -0700 Subject: [PATCH 10/11] Debug assert that the pool is empty on drop, per Alex's suggestion Also fix a couple scenarios where we could leak indices if allocating an index for a memory/table succeeded but then creating the memory/table itself failed. --- .../runtime/src/instance/allocator/pooling.rs | 43 ++++++++- .../allocator/pooling/index_allocator.rs | 13 +++ .../instance/allocator/pooling/memory_pool.rs | 95 +++++++++++-------- .../instance/allocator/pooling/stack_pool.rs | 5 + .../instance/allocator/pooling/table_pool.rs | 39 +++++--- 5 files changed, 135 insertions(+), 60 deletions(-) diff --git a/crates/runtime/src/instance/allocator/pooling.rs b/crates/runtime/src/instance/allocator/pooling.rs index 41a5579fc384..0f3bd25bd96f 100644 --- a/crates/runtime/src/instance/allocator/pooling.rs +++ b/crates/runtime/src/instance/allocator/pooling.rs @@ -202,8 +202,26 @@ pub struct PoolingInstanceAllocator { #[cfg(all(feature = "async", unix, not(miri)))] stacks: StackPool, + #[cfg(all(feature = "async", windows))] stack_size: usize, + #[cfg(all(feature = "async", windows))] + live_stacks: AtomicU64, +} + +impl Drop for PoolingInstanceAllocator { + fn drop(&mut self) { + debug_assert_eq!(self.live_component_instances.load(Ordering::Acquire), 0); + debug_assert_eq!(self.live_core_instances.load(Ordering::Acquire), 0); + + debug_assert!(self.memories.is_empty()); + debug_assert!(self.tables.is_empty()); + + #[cfg(all(feature = "async", unix, not(miri)))] + debug_assert!(self.stacks.is_empty()); + #[cfg(all(feature = "async", windows))] + debug_assert_eq!(self.live_stacks.load(Ordering::Acquire), 0); + } } impl PoolingInstanceAllocator { @@ -467,9 +485,25 @@ unsafe impl InstanceAllocatorImpl for PoolingInstanceAllocator { } // On windows, we don't use a stack pool as we use the native - // fiber implementation - let stack = wasmtime_fiber::FiberStack::new(self.stack_size)?; - Ok(stack) + // fiber implementation. We do still enforce the `total_stacks` + // limit, however. + + let old_count = self.live_stacks.fetch_add(1, Ordering::AcqRel); + if old_count >= u64::from(self.limits.total_stacks) { + self.live_stacks.fetch_sub(1, Ordering::AcqRel); + bail!( + "maximum concurrent fiber limit of {} reached", + self.limits.total_stacks + ); + } + + match wasmtime_fiber::FiberStack::new(self.stack_size) { + Ok(stack) => Ok(stack), + Err(e) => { + self.live_stacks.fetch_sub(1, Ordering::AcqRel); + Err(e) + } + } } else { compile_error!("not implemented"); } @@ -485,7 +519,8 @@ unsafe impl InstanceAllocatorImpl for PoolingInstanceAllocator { } else if #[cfg(unix)] { self.stacks.deallocate(stack); } else if #[cfg(windows)] { - // A no-op as we don't own the fiber stack on Windows + self.live_stacks.fetch_sub(1, Ordering::AcqRel); + // A no-op as we don't own the fiber stack on Windows. let _ = stack; } else { compile_error!("not implemented"); diff --git a/crates/runtime/src/instance/allocator/pooling/index_allocator.rs b/crates/runtime/src/instance/allocator/pooling/index_allocator.rs index c7acf5d97d6b..d4079680be89 100644 --- a/crates/runtime/src/instance/allocator/pooling/index_allocator.rs +++ b/crates/runtime/src/instance/allocator/pooling/index_allocator.rs @@ -30,6 +30,10 @@ impl SimpleIndexAllocator { SimpleIndexAllocator(ModuleAffinityIndexAllocator::new(capacity, 0)) } + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } + pub fn alloc(&self) -> Option { self.0.alloc(None) } @@ -163,6 +167,15 @@ impl ModuleAffinityIndexAllocator { })) } + /// Are zero slots in use right now? + pub fn is_empty(&self) -> bool { + let inner = self.0.lock().unwrap(); + !inner + .slot_state + .iter() + .any(|s| matches!(s, SlotState::Used(_))) + } + /// Allocate a new index from this allocator optionally using `id` as an /// affinity request if the allocation strategy supports it. /// diff --git a/crates/runtime/src/instance/allocator/pooling/memory_pool.rs b/crates/runtime/src/instance/allocator/pooling/memory_pool.rs index e215d59d5457..45038a21e860 100644 --- a/crates/runtime/src/instance/allocator/pooling/memory_pool.rs +++ b/crates/runtime/src/instance/allocator/pooling/memory_pool.rs @@ -199,6 +199,11 @@ impl MemoryPool { Ok(()) } + /// Are zero slots in use right now? + pub fn is_empty(&self) -> bool { + self.index_allocator.is_empty() + } + /// Allocate a single memory for the given instance allocation request. pub fn allocate( &self, @@ -222,50 +227,56 @@ impl MemoryPool { ) })?; - // Double-check that the runtime requirements of the memory are - // satisfied by the configuration of this pooling allocator. This - // should be returned as an error through `validate_memory_plans` - // but double-check here to be sure. - match memory_plan.style { - MemoryStyle::Static { bound } => { - let bound = bound * u64::from(WASM_PAGE_SIZE); - assert!(bound <= u64::try_from(self.memory_size).unwrap()); + match (|| { + // Double-check that the runtime requirements of the memory are + // satisfied by the configuration of this pooling allocator. This + // should be returned as an error through `validate_memory_plans` + // but double-check here to be sure. + match memory_plan.style { + MemoryStyle::Static { bound } => { + let bound = bound * u64::from(WASM_PAGE_SIZE); + assert!(bound <= u64::try_from(self.memory_size).unwrap()); + } + MemoryStyle::Dynamic { .. } => {} } - MemoryStyle::Dynamic { .. } => {} - } - let base_ptr = self.get_base(allocation_index); - let base_capacity = self.max_accessible; - - let mut slot = self.take_memory_image_slot(allocation_index); - let image = request.runtime_info.memory_image(memory_index)?; - let initial_size = memory_plan.memory.minimum * WASM_PAGE_SIZE as u64; - - // If instantiation fails, we can propagate the error - // upward and drop the slot. This will cause the Drop - // handler to attempt to map the range with PROT_NONE - // memory, to reserve the space while releasing any - // stale mappings. The next use of this slot will then - // create a new slot that will try to map over - // this, returning errors as well if the mapping - // errors persist. The unmap-on-drop is best effort; - // if it fails, then we can still soundly continue - // using the rest of the pool and allowing the rest of - // the process to continue, because we never perform a - // mmap that would leave an open space for someone - // else to come in and map something. - slot.instantiate(initial_size as usize, image, memory_plan)?; - - let memory = Memory::new_static( - memory_plan, - base_ptr, - base_capacity, - slot, - self.memory_and_guard_size, - unsafe { &mut *request.store.get().unwrap() }, - )?; - - Ok((allocation_index, memory)) + let base_ptr = self.get_base(allocation_index); + let base_capacity = self.max_accessible; + + let mut slot = self.take_memory_image_slot(allocation_index); + let image = request.runtime_info.memory_image(memory_index)?; + let initial_size = memory_plan.memory.minimum * WASM_PAGE_SIZE as u64; + + // If instantiation fails, we can propagate the error + // upward and drop the slot. This will cause the Drop + // handler to attempt to map the range with PROT_NONE + // memory, to reserve the space while releasing any + // stale mappings. The next use of this slot will then + // create a new slot that will try to map over + // this, returning errors as well if the mapping + // errors persist. The unmap-on-drop is best effort; + // if it fails, then we can still soundly continue + // using the rest of the pool and allowing the rest of + // the process to continue, because we never perform a + // mmap that would leave an open space for someone + // else to come in and map something. + slot.instantiate(initial_size as usize, image, memory_plan)?; + + Memory::new_static( + memory_plan, + base_ptr, + base_capacity, + slot, + self.memory_and_guard_size, + unsafe { &mut *request.store.get().unwrap() }, + ) + })() { + Ok(memory) => Ok((allocation_index, memory)), + Err(e) => { + self.index_allocator.free(SlotId(allocation_index.0)); + Err(e) + } + } } /// Deallocate a previously-allocated memory. diff --git a/crates/runtime/src/instance/allocator/pooling/stack_pool.rs b/crates/runtime/src/instance/allocator/pooling/stack_pool.rs index 651e212646a7..8265cfcf94a7 100644 --- a/crates/runtime/src/instance/allocator/pooling/stack_pool.rs +++ b/crates/runtime/src/instance/allocator/pooling/stack_pool.rs @@ -74,6 +74,11 @@ impl StackPool { }) } + /// Are there zero slots in use right now? + pub fn is_empty(&self) -> bool { + self.index_allocator.is_empty() + } + /// Allocate a new fiber. pub fn allocate(&self) -> Result { if self.stack_size == 0 { diff --git a/crates/runtime/src/instance/allocator/pooling/table_pool.rs b/crates/runtime/src/instance/allocator/pooling/table_pool.rs index a5607c46c4de..81938cb0491d 100644 --- a/crates/runtime/src/instance/allocator/pooling/table_pool.rs +++ b/crates/runtime/src/instance/allocator/pooling/table_pool.rs @@ -91,6 +91,11 @@ impl TablePool { Ok(()) } + /// Are there zero slots in use right now? + pub fn is_empty(&self) -> bool { + self.index_allocator.is_empty() + } + /// Get the base pointer of the given table allocation. fn get(&self, table_index: TableAllocationIndex) -> *mut u8 { assert!(table_index.index() < self.max_total_tables); @@ -120,20 +125,26 @@ impl TablePool { ) })?; - let base = self.get(allocation_index); - - commit_table_pages( - base as *mut u8, - self.table_elements * mem::size_of::<*mut u8>(), - )?; - - let table = Table::new_static( - table_plan, - unsafe { std::slice::from_raw_parts_mut(base.cast(), self.table_elements) }, - unsafe { &mut *request.store.get().unwrap() }, - )?; - - Ok((allocation_index, table)) + match (|| { + let base = self.get(allocation_index); + + commit_table_pages( + base as *mut u8, + self.table_elements * mem::size_of::<*mut u8>(), + )?; + + Table::new_static( + table_plan, + unsafe { std::slice::from_raw_parts_mut(base.cast(), self.table_elements) }, + unsafe { &mut *request.store.get().unwrap() }, + ) + })() { + Ok(table) => Ok((allocation_index, table)), + Err(e) => { + self.index_allocator.free(SlotId(allocation_index.0)); + Err(e) + } + } } /// Deallocate a previously-allocated table. From f37b84b9450c102116e95148ec66c49d3f621742 Mon Sep 17 00:00:00 2001 From: Nick Fitzgerald Date: Fri, 18 Aug 2023 14:00:31 -0700 Subject: [PATCH 11/11] Fix windows compile errors --- crates/runtime/src/instance/allocator/pooling.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/crates/runtime/src/instance/allocator/pooling.rs b/crates/runtime/src/instance/allocator/pooling.rs index 0f3bd25bd96f..214e66c8c5f5 100644 --- a/crates/runtime/src/instance/allocator/pooling.rs +++ b/crates/runtime/src/instance/allocator/pooling.rs @@ -237,6 +237,8 @@ impl PoolingInstanceAllocator { stacks: StackPool::new(config)?, #[cfg(all(feature = "async", windows))] stack_size: config.stack_size, + #[cfg(all(feature = "async", windows))] + live_stacks: AtomicU64::new(0), }) } @@ -501,7 +503,7 @@ unsafe impl InstanceAllocatorImpl for PoolingInstanceAllocator { Ok(stack) => Ok(stack), Err(e) => { self.live_stacks.fetch_sub(1, Ordering::AcqRel); - Err(e) + Err(anyhow::Error::from(e)) } } } else {