diff --git a/crates/runtime/src/instance.rs b/crates/runtime/src/instance.rs index 9c56dfb22229..d2afb4d871b4 100644 --- a/crates/runtime/src/instance.rs +++ b/crates/runtime/src/instance.rs @@ -11,7 +11,7 @@ use crate::vmcontext::{ VMCallerCheckedAnyfunc, VMContext, VMFunctionImport, VMGlobalDefinition, VMGlobalImport, VMInterrupts, VMMemoryDefinition, VMMemoryImport, VMTableDefinition, VMTableImport, }; -use crate::{ExportFunction, ExportGlobal, ExportMemory, ExportTable, Store}; +use crate::{CompiledModuleId, ExportFunction, ExportGlobal, ExportMemory, ExportTable, Store}; use anyhow::Error; use memoffset::offset_of; use more_asserts::assert_lt; @@ -54,6 +54,9 @@ pub(crate) struct Instance { /// The `Module` this `Instance` was instantiated from. module: Arc, + /// The unique ID for the `Module` this `Instance` was instantiated from. + unique_id: Option, + /// Offsets in the `vmctx` region, precomputed from the `module` above. offsets: VMOffsets, @@ -100,6 +103,7 @@ impl Instance { /// Helper for allocators; not a public API. pub(crate) fn create_raw( module: &Arc, + unique_id: Option, wasm_data: &'static [u8], memories: PrimaryMap, tables: PrimaryMap, @@ -107,6 +111,7 @@ impl Instance { ) -> Instance { Instance { module: module.clone(), + unique_id, offsets: VMOffsets::new(HostPtr, &module), memories, tables, diff --git a/crates/runtime/src/instance/allocator.rs b/crates/runtime/src/instance/allocator.rs index c30f0dfd9560..a1c50fa8c549 100644 --- a/crates/runtime/src/instance/allocator.rs +++ b/crates/runtime/src/instance/allocator.rs @@ -7,7 +7,7 @@ use crate::vmcontext::{ VMBuiltinFunctionsArray, VMCallerCheckedAnyfunc, VMGlobalDefinition, VMSharedSignatureIndex, }; use crate::ModuleMemFds; -use crate::Store; +use crate::{CompiledModuleId, Store}; use anyhow::Result; use std::alloc; use std::any::Any; @@ -80,6 +80,9 @@ pub struct InstanceAllocationRequest<'a> { /// The module being instantiated. pub module: Arc, + /// The unique ID of the module being allocated within this engine. + pub unique_id: Option, + /// The base address of where JIT functions are located. pub image_base: usize, @@ -760,8 +763,14 @@ unsafe impl InstanceAllocator for OnDemandInstanceAllocator { let host_state = std::mem::replace(&mut req.host_state, Box::new(())); let mut handle = { - let instance = - Instance::create_raw(&req.module, &*req.wasm_data, memories, tables, host_state); + let instance = Instance::create_raw( + &req.module, + req.unique_id, + &*req.wasm_data, + memories, + tables, + host_state, + ); let layout = instance.alloc_layout(); let instance_ptr = alloc::alloc(layout) as *mut Instance; if instance_ptr.is_null() { diff --git a/crates/runtime/src/instance/allocator/pooling.rs b/crates/runtime/src/instance/allocator/pooling.rs index afa1aee56ffa..46a03a4a798c 100644 --- a/crates/runtime/src/instance/allocator/pooling.rs +++ b/crates/runtime/src/instance/allocator/pooling.rs @@ -16,7 +16,6 @@ use crate::memfd::MemoryMemFd; use crate::{instance::Instance, Memory, Mmap, ModuleMemFds, Table}; use anyhow::{anyhow, bail, Context, Result}; use libc::c_void; -use rand::Rng; #[cfg(feature = "memfd-allocator")] use rustix::fd::AsRawFd; use std::convert::TryFrom; @@ -30,6 +29,9 @@ use wasmtime_environ::{ WASM_PAGE_SIZE, }; +mod index_allocator; +use index_allocator::PoolingAllocationState; + cfg_if::cfg_if! { if #[cfg(windows)] { mod windows; @@ -255,20 +257,19 @@ pub enum PoolingAllocationStrategy { NextAvailable, /// Allocate from a random available instance. Random, + /// Try to allocate an instance slot that was previously used for + /// the same module, potentially enabling faster instantiation by + /// reusing e.g. memory mappings. + ReuseAffinity, } -impl PoolingAllocationStrategy { - fn next(&self, free_count: usize) -> usize { - debug_assert!(free_count > 0); - - match self { - Self::NextAvailable => free_count - 1, - Self::Random => rand::thread_rng().gen_range(0..free_count), - } +impl Default for PoolingAllocationStrategy { + #[cfg(feature = "memfd-allocator")] + fn default() -> Self { + Self::ReuseAffinity } -} -impl Default for PoolingAllocationStrategy { + #[cfg(not(feature = "memfd-allocator"))] fn default() -> Self { Self::NextAvailable } @@ -288,13 +289,14 @@ struct InstancePool { mapping: Mmap, instance_size: usize, max_instances: usize, - free_list: Mutex>, + index_allocator: Mutex, memories: MemoryPool, tables: TablePool, } impl InstancePool { fn new( + strategy: PoolingAllocationStrategy, module_limits: &ModuleLimits, instance_limits: &InstanceLimits, tunables: &Tunables, @@ -335,7 +337,7 @@ impl InstancePool { mapping, instance_size, max_instances, - free_list: Mutex::new((0..max_instances).collect()), + index_allocator: Mutex::new(PoolingAllocationState::new(strategy, max_instances)), memories: MemoryPool::new(module_limits, instance_limits, tunables)?, tables: TablePool::new(module_limits, instance_limits)?, }; @@ -356,6 +358,7 @@ impl InstancePool { let host_state = std::mem::replace(&mut req.host_state, Box::new(())); let instance_data = Instance::create_raw( &req.module, + req.unique_id, &*req.wasm_data, PrimaryMap::default(), PrimaryMap::default(), @@ -367,6 +370,7 @@ impl InstancePool { // chosen slot before we do anything else with it. (This is // paired with a `drop_in_place` in deallocate below.) let instance = self.instance(index); + std::ptr::write(instance as _, instance_data); // set_instance_memories and _tables will need the store before we can completely @@ -398,16 +402,14 @@ impl InstancePool { fn allocate( &self, - strategy: PoolingAllocationStrategy, req: InstanceAllocationRequest, ) -> Result { let index = { - let mut free_list = self.free_list.lock().unwrap(); - if free_list.is_empty() { + let mut alloc = self.index_allocator.lock().unwrap(); + if alloc.is_empty() { return Err(InstantiationError::Limit(self.max_instances as u32)); } - let free_index = strategy.next(free_list.len()); - free_list.swap_remove(free_index) + alloc.alloc(req.unique_id) }; unsafe { @@ -433,6 +435,7 @@ impl InstancePool { debug_assert!(index < self.max_instances); let instance = unsafe { &mut *handle.instance }; + let unique_id = instance.unique_id; // Decommit any linear memories that were used for ((def_mem_idx, memory), base) in @@ -497,7 +500,7 @@ impl InstancePool { // touched again until we write a fresh Instance in-place with // std::ptr::write in allocate() above. - self.free_list.lock().unwrap().push(index); + self.index_allocator.lock().unwrap().free(index, unique_id); } fn set_instance_memories( @@ -1093,7 +1096,7 @@ struct StackPool { stack_size: usize, max_instances: usize, page_size: usize, - free_list: Mutex>, + index_allocator: Mutex, } #[cfg(all(feature = "async", unix))] @@ -1136,25 +1139,29 @@ impl StackPool { stack_size, max_instances, page_size, - free_list: Mutex::new((0..max_instances).collect()), + // We always use a `NextAvailable` strategy for stack + // allocation. We don't want or need an affinity policy + // here: stacks do not benefit from being allocated to the + // same compiled module with the same image (they always + // start zeroed just the same for everyone). + index_allocator: Mutex::new(PoolingAllocationState::new( + PoolingAllocationStrategy::NextAvailable, + max_instances, + )), }) } - fn allocate( - &self, - strategy: PoolingAllocationStrategy, - ) -> Result { + fn allocate(&self) -> Result { if self.stack_size == 0 { return Err(FiberStackError::NotSupported); } let index = { - let mut free_list = self.free_list.lock().unwrap(); - if free_list.is_empty() { + let mut alloc = self.index_allocator.lock().unwrap(); + if alloc.is_empty() { return Err(FiberStackError::Limit(self.max_instances as u32)); } - let free_index = strategy.next(free_list.len()); - free_list.swap_remove(free_index) + alloc.alloc(None) }; debug_assert!(index < self.max_instances); @@ -1200,7 +1207,7 @@ impl StackPool { decommit_stack_pages(bottom_of_stack as _, stack_size).unwrap(); - self.free_list.lock().unwrap().push(index); + self.index_allocator.lock().unwrap().free(index, None); } } @@ -1211,7 +1218,6 @@ impl StackPool { /// Note: the resource pools are manually dropped so that the fault handler terminates correctly. #[derive(Debug)] pub struct PoolingInstanceAllocator { - strategy: PoolingAllocationStrategy, module_limits: ModuleLimits, // This is manually drop so that the pools unmap their memory before the page fault handler drops. instances: mem::ManuallyDrop, @@ -1236,7 +1242,7 @@ impl PoolingInstanceAllocator { bail!("the instance count limit cannot be zero"); } - let instances = InstancePool::new(&module_limits, &instance_limits, tunables)?; + let instances = InstancePool::new(strategy, &module_limits, &instance_limits, tunables)?; #[cfg(all(feature = "uffd", target_os = "linux"))] let _fault_handler = imp::PageFaultHandler::new(&instances)?; @@ -1244,7 +1250,6 @@ impl PoolingInstanceAllocator { drop(stack_size); // suppress unused warnings w/o async feature Ok(Self { - strategy, module_limits, instances: mem::ManuallyDrop::new(instances), #[cfg(all(feature = "async", unix))] @@ -1283,7 +1288,7 @@ unsafe impl InstanceAllocator for PoolingInstanceAllocator { &self, req: InstanceAllocationRequest, ) -> Result { - self.instances.allocate(self.strategy, req) + self.instances.allocate(req) } unsafe fn initialize( @@ -1330,7 +1335,7 @@ unsafe impl InstanceAllocator for PoolingInstanceAllocator { #[cfg(all(feature = "async", unix))] fn allocate_fiber_stack(&self) -> Result { - self.stacks.allocate(self.strategy) + self.stacks.allocate() } #[cfg(all(feature = "async", unix))] @@ -1650,21 +1655,6 @@ mod test { ); } - #[test] - fn test_next_available_allocation_strategy() { - let strat = PoolingAllocationStrategy::NextAvailable; - assert_eq!(strat.next(10), 9); - assert_eq!(strat.next(5), 4); - assert_eq!(strat.next(1), 0); - } - - #[test] - fn test_random_allocation_strategy() { - let strat = PoolingAllocationStrategy::Random; - assert!(strat.next(100) < 100); - assert_eq!(strat.next(1), 0); - } - #[cfg(target_pointer_width = "64")] #[test] fn test_instance_pool() -> Result<()> { @@ -1684,6 +1674,7 @@ mod test { let instance_limits = InstanceLimits { count: 3 }; let instances = InstancePool::new( + PoolingAllocationStrategy::NextAvailable, &module_limits, &instance_limits, &Tunables { @@ -1697,7 +1688,10 @@ mod test { assert_eq!(instances.instance_size, region::page::size()); assert_eq!(instances.max_instances, 3); - assert_eq!(&*instances.free_list.lock().unwrap(), &[0, 1, 2]); + assert_eq!( + instances.index_allocator.lock().unwrap().testing_freelist(), + &[0, 1, 2] + ); let mut handles = Vec::new(); let module = Arc::new(Module::default()); @@ -1706,50 +1700,49 @@ mod test { for _ in (0..3).rev() { handles.push( instances - .allocate( - PoolingAllocationStrategy::NextAvailable, - InstanceAllocationRequest { - module: module.clone(), - image_base: 0, - functions, - imports: Imports { - functions: &[], - tables: &[], - memories: &[], - globals: &[], - }, - shared_signatures: VMSharedSignatureIndex::default().into(), - host_state: Box::new(()), - store: StorePtr::empty(), - wasm_data: &[], - memfds: None, + .allocate(InstanceAllocationRequest { + module: module.clone(), + unique_id: None, + image_base: 0, + functions, + imports: Imports { + functions: &[], + tables: &[], + memories: &[], + globals: &[], }, - ) + shared_signatures: VMSharedSignatureIndex::default().into(), + host_state: Box::new(()), + store: StorePtr::empty(), + wasm_data: &[], + memfds: None, + }) .expect("allocation should succeed"), ); } - assert_eq!(&*instances.free_list.lock().unwrap(), &[]); + assert_eq!( + instances.index_allocator.lock().unwrap().testing_freelist(), + &[] + ); - match instances.allocate( - PoolingAllocationStrategy::NextAvailable, - InstanceAllocationRequest { - module: module.clone(), - functions, - image_base: 0, - imports: Imports { - functions: &[], - tables: &[], - memories: &[], - globals: &[], - }, - shared_signatures: VMSharedSignatureIndex::default().into(), - host_state: Box::new(()), - store: StorePtr::empty(), - wasm_data: &[], - memfds: None, + match instances.allocate(InstanceAllocationRequest { + module: module.clone(), + unique_id: None, + functions, + image_base: 0, + imports: Imports { + functions: &[], + tables: &[], + memories: &[], + globals: &[], }, - ) { + shared_signatures: VMSharedSignatureIndex::default().into(), + host_state: Box::new(()), + store: StorePtr::empty(), + wasm_data: &[], + memfds: None, + }) { Err(InstantiationError::Limit(3)) => {} _ => panic!("unexpected error"), }; @@ -1758,7 +1751,10 @@ mod test { instances.deallocate(&handle); } - assert_eq!(&*instances.free_list.lock().unwrap(), &[2, 1, 0]); + assert_eq!( + instances.index_allocator.lock().unwrap().testing_freelist(), + &[2, 1, 0] + ); Ok(()) } @@ -1868,7 +1864,7 @@ mod test { assert_eq!(pool.page_size, native_page_size); assert_eq!( - &*pool.free_list.lock().unwrap(), + pool.index_allocator.lock().unwrap().testing_freelist(), &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], ); @@ -1876,9 +1872,7 @@ mod test { let mut stacks = Vec::new(); for i in (0..10).rev() { - let stack = pool - .allocate(PoolingAllocationStrategy::NextAvailable) - .expect("allocation should succeed"); + let stack = pool.allocate().expect("allocation should succeed"); assert_eq!( ((stack.top().unwrap() as usize - base) / pool.stack_size) - 1, i @@ -1886,12 +1880,9 @@ mod test { stacks.push(stack); } - assert_eq!(&*pool.free_list.lock().unwrap(), &[]); + assert_eq!(pool.index_allocator.lock().unwrap().testing_freelist(), &[]); - match pool - .allocate(PoolingAllocationStrategy::NextAvailable) - .unwrap_err() - { + match pool.allocate().unwrap_err() { FiberStackError::Limit(10) => {} _ => panic!("unexpected error"), }; @@ -1901,7 +1892,7 @@ mod test { } assert_eq!( - &*pool.free_list.lock().unwrap(), + pool.index_allocator.lock().unwrap().testing_freelist(), &[9, 8, 7, 6, 5, 4, 3, 2, 1, 0], ); diff --git a/crates/runtime/src/instance/allocator/pooling/index_allocator.rs b/crates/runtime/src/instance/allocator/pooling/index_allocator.rs new file mode 100644 index 000000000000..ee4485b8edb9 --- /dev/null +++ b/crates/runtime/src/instance/allocator/pooling/index_allocator.rs @@ -0,0 +1,408 @@ +//! Index/slot allocator policies for the pooling allocator. + +use super::PoolingAllocationStrategy; +use crate::CompiledModuleId; +use rand::Rng; +use std::collections::HashMap; + +#[derive(Clone, Debug)] +pub(crate) enum PoolingAllocationState { + NextAvailable(Vec), + Random(Vec), + /// Reuse-affinity policy state. + /// + /// The data structures here deserve a little explanation: + /// + /// - free_list: this is a vec of slot indices that are free, no + /// matter their affinities. + /// - per_module: this is a hashmap of vecs of slot indices that + /// are free, with affinity for particular module IDs. A slot may + /// appear in zero or one of these lists. + /// - slot_state: indicates what state each slot is in: allocated + /// (Taken), only in free_list (Empty), or in free_list and a + /// per_module list (Affinity). + /// + /// The slot state tracks a slot's index in the global and + /// per-module freelists, so it can be efficiently removed from + /// both. We take some care to keep these up-to-date as well. + /// + /// On allocation, we first try to find a slot with affinity for + /// the given module ID, if any. If not, we pick a random slot + /// ID. This random choice is unbiased across all free slots. + ReuseAffinity { + // Free-list of all slots. We use this to pick a victim when + // we don't have an appropriate slot with the preferred + // affinity. + free_list: Vec, + // Invariant: any module ID in this hashmap must have a + // non-empty list of free slots (otherwise we remove it). + per_module: HashMap>, + // The state of any given slot. Records indices in the above + // list (empty) or two lists (with affinity), and these + // indices are kept up-to-date to allow fast removal. + slot_state: Vec, + }, +} + +#[derive(Clone, Debug)] +pub(crate) enum SlotState { + Taken, + Empty { + /// Index in the global free list. Invariant: + /// free_list[slot_state[i].free_list_index] == i. + free_list_index: usize, + }, + Affinity { + module: CompiledModuleId, + /// Index in the global free list. Invariant: + /// free_list[slot_state[i].free_list_index] == i. + free_list_index: usize, + /// Index in a per-module free list. Invariant: + /// per_module[slot_state[i].module][slot_state[i].per_module_index] + /// == i. + per_module_index: usize, + }, +} + +impl SlotState { + /// Get the index of this slot in the global free list. + fn free_list_index(&self) -> usize { + match self { + &Self::Empty { free_list_index } + | &Self::Affinity { + free_list_index, .. + } => free_list_index, + _ => unreachable!(), + } + } + + /// Update the index of this slot in the global free list. + fn update_free_list_index(&mut self, index: usize) { + match self { + &mut Self::Empty { + ref mut free_list_index, + } + | &mut Self::Affinity { + ref mut free_list_index, + .. + } => { + *free_list_index = index; + } + _ => panic!("Taken in free list"), + } + } + + /// Get the index of this slot in its per-module free list. + fn per_module_index(&self) -> usize { + match self { + &Self::Affinity { + per_module_index, .. + } => per_module_index, + _ => unreachable!(), + } + } + + /// Update the index of this slot in its per-module free list. + fn update_per_module_index(&mut self, index: usize) { + match self { + &mut Self::Affinity { + ref mut per_module_index, + .. + } => { + *per_module_index = index; + } + _ => panic!("Taken in per-module free list"), + } + } +} + +impl PoolingAllocationState { + /// Create the default state for this strategy. + pub(crate) fn new(strategy: PoolingAllocationStrategy, max_instances: usize) -> Self { + let ids = (0..max_instances).collect::>(); + match strategy { + PoolingAllocationStrategy::NextAvailable => PoolingAllocationState::NextAvailable(ids), + PoolingAllocationStrategy::Random => PoolingAllocationState::Random(ids), + PoolingAllocationStrategy::ReuseAffinity => PoolingAllocationState::ReuseAffinity { + free_list: ids, + per_module: HashMap::new(), + slot_state: (0..max_instances) + .map(|i| SlotState::Empty { free_list_index: i }) + .collect(), + }, + } + } + + /// Are any slots left, or is this allocator empty? + pub(crate) fn is_empty(&self) -> bool { + match self { + &PoolingAllocationState::NextAvailable(ref free_list) + | &PoolingAllocationState::Random(ref free_list) => free_list.is_empty(), + &PoolingAllocationState::ReuseAffinity { ref free_list, .. } => free_list.is_empty(), + } + } + + /// Internal: remove a slot-index from the global free list. + fn remove_free_list_item( + slot_state: &mut Vec, + free_list: &mut Vec, + index: usize, + ) { + let free_list_index = slot_state[index].free_list_index(); + assert_eq!(index, free_list.swap_remove(free_list_index)); + if free_list_index < free_list.len() { + let replaced = free_list[free_list_index]; + slot_state[replaced].update_free_list_index(free_list_index); + } + } + + /// Internal: remove a slot-index from a per-module free list. + fn remove_module_free_list_item( + slot_state: &mut Vec, + per_module: &mut HashMap>, + id: CompiledModuleId, + index: usize, + ) { + let per_module_list = per_module.get_mut(&id).unwrap(); + let per_module_index = slot_state[index].per_module_index(); + assert_eq!(index, per_module_list.swap_remove(per_module_index)); + if per_module_index < per_module_list.len() { + let replaced = per_module_list[per_module_index]; + slot_state[replaced].update_per_module_index(per_module_index); + } + if per_module_list.is_empty() { + per_module.remove(&id); + } + } + + /// Allocate a new slot. + pub(crate) fn alloc(&mut self, id: Option) -> usize { + match self { + &mut PoolingAllocationState::NextAvailable(ref mut free_list) => { + debug_assert!(free_list.len() > 0); + free_list.pop().unwrap() + } + &mut PoolingAllocationState::Random(ref mut free_list) => { + debug_assert!(free_list.len() > 0); + let id = rand::thread_rng().gen_range(0..free_list.len()); + free_list.swap_remove(id) + } + &mut PoolingAllocationState::ReuseAffinity { + ref mut free_list, + ref mut per_module, + ref mut slot_state, + .. + } => { + if let Some(this_module) = id.and_then(|id| per_module.get_mut(&id)) { + // There is a freelist of slots with affinity for + // the requested module-ID. Pick the last one; any + // will do, no need for randomness here. + assert!(!this_module.is_empty()); + let new_id = this_module.pop().expect("List should never be empty"); + if this_module.is_empty() { + per_module.remove(&id.unwrap()); + } + // Make sure to remove from the global + // freelist. We already removed from the + // per-module list above. + Self::remove_free_list_item(slot_state, free_list, new_id); + slot_state[new_id] = SlotState::Taken; + new_id + } else { + // Pick a random free slot ID. Note that we do + // this, rather than pick a victim module first, + // to maintain an unbiased stealing distribution: + // we want the likelihood of our taking a slot + // from some other module's freelist to be + // proportional to that module's freelist + // length. Or in other words, every *slot* should + // be equally likely to be stolen. The + // alternative, where we pick the victim module + // freelist first, means that either a module with + // an affinity freelist of one slot has the same + // chances of losing that slot as one with a + // hundred slots; or else we need a weighted + // random choice among modules, which is just as + // complex as this process. + // + // We don't bother picking an empty slot (no + // established affinity) before a random slot, + // because this is more complex, and in the steady + // state, all slots will see at least one + // instantiation very quickly, so there will never + // (past an initial phase) be a slot with no + // affinity. + let free_list_index = rand::thread_rng().gen_range(0..free_list.len()); + let new_id = free_list[free_list_index]; + // Remove from both the global freelist and + // per-module freelist, if any. + Self::remove_free_list_item(slot_state, free_list, new_id); + if let &SlotState::Affinity { module, .. } = &slot_state[new_id] { + Self::remove_module_free_list_item(slot_state, per_module, module, new_id); + } + slot_state[new_id] = SlotState::Taken; + + new_id + } + } + } + } + + pub(crate) fn free(&mut self, index: usize, id: Option) { + match self { + &mut PoolingAllocationState::NextAvailable(ref mut free_list) + | &mut PoolingAllocationState::Random(ref mut free_list) => { + free_list.push(index); + } + &mut PoolingAllocationState::ReuseAffinity { + ref mut per_module, + ref mut free_list, + ref mut slot_state, + } => { + let free_list_index = free_list.len(); + free_list.push(index); + if let Some(id) = id { + let per_module_list = per_module.entry(id).or_insert_with(|| vec![]); + let per_module_index = per_module_list.len(); + per_module_list.push(index); + slot_state[index] = SlotState::Affinity { + module: id, + free_list_index, + per_module_index, + }; + } else { + slot_state[index] = SlotState::Empty { free_list_index }; + } + } + } + } + + /// 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)] + pub(crate) fn testing_freelist(&self) -> &[usize] { + match self { + &PoolingAllocationState::NextAvailable(ref free_list) + | &PoolingAllocationState::Random(ref free_list) => &free_list[..], + _ => panic!("Wrong kind of state"), + } + } +} + +#[cfg(test)] +mod test { + use super::PoolingAllocationState; + use crate::CompiledModuleIdAllocator; + use crate::PoolingAllocationStrategy; + + #[test] + fn test_next_available_allocation_strategy() { + let strat = PoolingAllocationStrategy::NextAvailable; + let mut state = PoolingAllocationState::new(strat, 10); + assert_eq!(state.alloc(None), 9); + let mut state = PoolingAllocationState::new(strat, 5); + assert_eq!(state.alloc(None), 4); + let mut state = PoolingAllocationState::new(strat, 1); + assert_eq!(state.alloc(None), 0); + } + + #[test] + fn test_random_allocation_strategy() { + let strat = PoolingAllocationStrategy::Random; + let mut state = PoolingAllocationState::new(strat, 100); + assert!(state.alloc(None) < 100); + let mut state = PoolingAllocationState::new(strat, 1); + assert_eq!(state.alloc(None), 0); + } + + #[test] + fn test_affinity_allocation_strategy() { + let strat = PoolingAllocationStrategy::ReuseAffinity; + let id_alloc = CompiledModuleIdAllocator::new(); + let id1 = id_alloc.alloc(); + let id2 = id_alloc.alloc(); + let mut state = PoolingAllocationState::new(strat, 100); + + let index1 = state.alloc(Some(id1)); + assert!(index1 < 100); + let index2 = state.alloc(Some(id2)); + assert!(index2 < 100); + assert_ne!(index1, index2); + + state.free(index1, Some(id1)); + let index3 = state.alloc(Some(id1)); + assert_eq!(index3, index1); + state.free(index3, Some(id1)); + + state.free(index2, Some(id2)); + + // Now there is 1 free instance for id2 and 1 free instance + // for id1, and 98 empty. Allocate 100 for id2. The first + // should be equal to the one we know was previously used for + // id2. The next 99 are arbitrary. + + let mut indices = vec![]; + for _ in 0..100 { + assert!(!state.is_empty()); + indices.push(state.alloc(Some(id2))); + } + assert!(state.is_empty()); + assert_eq!(indices[0], index2); + + for i in indices { + state.free(i, Some(id2)); + } + + // Allocate an index we know previously had an instance but + // now does not (list ran empty). + let index = state.alloc(Some(id1)); + state.free(index, Some(id1)); + } + + #[test] + fn test_affinity_allocation_strategy_random() { + use rand::Rng; + let mut rng = rand::thread_rng(); + + let strat = PoolingAllocationStrategy::ReuseAffinity; + let id_alloc = CompiledModuleIdAllocator::new(); + let ids = std::iter::repeat_with(|| id_alloc.alloc()) + .take(10) + .collect::>(); + let mut state = PoolingAllocationState::new(strat, 1000); + let mut allocated = vec![]; + let mut last_id = vec![None; 1000]; + + let mut hits = 0; + for _ in 0..100_000 { + if !allocated.is_empty() && (state.is_empty() || rng.gen_bool(0.5)) { + let i = rng.gen_range(0..allocated.len()); + let (to_free_idx, to_free_id) = allocated.swap_remove(i); + let to_free_id = if rng.gen_bool(0.1) { + None + } else { + Some(to_free_id) + }; + state.free(to_free_idx, to_free_id); + } else { + assert!(!state.is_empty()); + let id = ids[rng.gen_range(0..ids.len())]; + let index = state.alloc(Some(id)); + if last_id[index] == Some(id) { + hits += 1; + } + last_id[index] = Some(id); + allocated.push((index, id)); + } + } + + // 10% reuse would be random chance (because we have 10 module + // IDs). Check for at least double that to ensure some sort of + // affinity is occurring. + assert!( + hits > 20000, + "expected at least 20000 (20%) ID-reuses but got {}", + hits + ); + } +} diff --git a/crates/runtime/src/instance/allocator/pooling/uffd.rs b/crates/runtime/src/instance/allocator/pooling/uffd.rs index 87dd9a0c57d5..be16ca2db1ec 100644 --- a/crates/runtime/src/instance/allocator/pooling/uffd.rs +++ b/crates/runtime/src/instance/allocator/pooling/uffd.rs @@ -466,8 +466,13 @@ mod test { ..Tunables::default() }; - let instances = InstancePool::new(&module_limits, &instance_limits, &tunables) - .expect("should allocate"); + let instances = InstancePool::new( + PoolingAllocationStrategy::Random, + &module_limits, + &instance_limits, + &tunables, + ) + .expect("should allocate"); let locator = FaultLocator::new(&instances); @@ -573,25 +578,23 @@ mod test { for _ in 0..instances.max_instances { handles.push( instances - .allocate( - PoolingAllocationStrategy::Random, - InstanceAllocationRequest { - module: module.clone(), - memfds: None, - image_base: 0, - functions, - imports: Imports { - functions: &[], - tables: &[], - memories: &[], - globals: &[], - }, - shared_signatures: VMSharedSignatureIndex::default().into(), - host_state: Box::new(()), - store: StorePtr::new(&mut mock_store), - wasm_data: &[], + .allocate(InstanceAllocationRequest { + module: module.clone(), + memfds: None, + unique_id: None, + image_base: 0, + functions, + imports: Imports { + functions: &[], + tables: &[], + memories: &[], + globals: &[], }, - ) + shared_signatures: VMSharedSignatureIndex::default().into(), + host_state: Box::new(()), + store: StorePtr::new(&mut mock_store), + wasm_data: &[], + }) .expect("instance should allocate"), ); } diff --git a/crates/runtime/src/module_id.rs b/crates/runtime/src/module_id.rs index 481a63e0bd3a..3c5ed7adf193 100644 --- a/crates/runtime/src/module_id.rs +++ b/crates/runtime/src/module_id.rs @@ -1,11 +1,14 @@ //! Unique IDs for modules in the runtime. -use std::sync::atomic::{AtomicU64, Ordering}; +use std::{ + num::NonZeroU64, + sync::atomic::{AtomicU64, Ordering}, +}; /// A unique identifier (within an engine or similar) for a compiled /// module. #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct CompiledModuleId(u64); +pub struct CompiledModuleId(NonZeroU64); /// An allocator for compiled module IDs. pub struct CompiledModuleIdAllocator { @@ -22,7 +25,19 @@ impl CompiledModuleIdAllocator { /// Allocate a new ID. pub fn alloc(&self) -> CompiledModuleId { + // Note: why is `Relaxed` OK here? + // + // The only requirement we have is that IDs are unique. We + // don't care how one module's ID compares to another, i.e., + // what order they come in. `Relaxed` means that this + // `fetch_add` operation does not have any particular + // synchronization (ordering) with respect to any other memory + // access in the program. However, `fetch_add` is always + // atomic with respect to other accesses to this variable + // (`self.next`). So we will always hand out separate, unique + // IDs correctly, just in some possibly arbitrary order (which + // is fine). let id = self.next.fetch_add(1, Ordering::Relaxed); - CompiledModuleId(id) + CompiledModuleId(NonZeroU64::new(id).unwrap()) } } diff --git a/crates/wasmtime/src/config/pooling.rs b/crates/wasmtime/src/config/pooling.rs index a2cbc470017e..6f8f9ae62b5f 100644 --- a/crates/wasmtime/src/config/pooling.rs +++ b/crates/wasmtime/src/config/pooling.rs @@ -249,6 +249,10 @@ pub enum PoolingAllocationStrategy { NextAvailable, /// Allocate from a random available instance. Random, + /// Try to allocate an instance slot that was previously used for + /// the same module, potentially enabling faster instantiation by + /// reusing e.g. memory mappings. + ReuseAffinity, } impl Default for PoolingAllocationStrategy { @@ -256,6 +260,7 @@ impl Default for PoolingAllocationStrategy { match wasmtime_runtime::PoolingAllocationStrategy::default() { wasmtime_runtime::PoolingAllocationStrategy::NextAvailable => Self::NextAvailable, wasmtime_runtime::PoolingAllocationStrategy::Random => Self::Random, + wasmtime_runtime::PoolingAllocationStrategy::ReuseAffinity => Self::ReuseAffinity, } } } @@ -268,6 +273,7 @@ impl Into for PoolingAllocationStra match self { Self::NextAvailable => wasmtime_runtime::PoolingAllocationStrategy::NextAvailable, Self::Random => wasmtime_runtime::PoolingAllocationStrategy::Random, + Self::ReuseAffinity => wasmtime_runtime::PoolingAllocationStrategy::ReuseAffinity, } } } diff --git a/crates/wasmtime/src/instance.rs b/crates/wasmtime/src/instance.rs index 7f5b5e823df7..99687621026a 100644 --- a/crates/wasmtime/src/instance.rs +++ b/crates/wasmtime/src/instance.rs @@ -707,6 +707,7 @@ impl<'a> Instantiator<'a> { .allocator() .allocate(InstanceAllocationRequest { module: compiled_module.module().clone(), + unique_id: Some(compiled_module.unique_id()), memfds: self.cur.module.memfds().clone(), image_base: compiled_module.code().as_ptr() as usize, functions: compiled_module.functions(), diff --git a/crates/wasmtime/src/store.rs b/crates/wasmtime/src/store.rs index 362fb5984843..56fd4fc3540a 100644 --- a/crates/wasmtime/src/store.rs +++ b/crates/wasmtime/src/store.rs @@ -421,6 +421,7 @@ impl Store { shared_signatures: None.into(), imports: Default::default(), module: Arc::new(wasmtime_environ::Module::default()), + unique_id: None, memfds: None, store: StorePtr::empty(), wasm_data: &[], diff --git a/crates/wasmtime/src/trampoline.rs b/crates/wasmtime/src/trampoline.rs index 790cbf9ef991..02e0b51c8130 100644 --- a/crates/wasmtime/src/trampoline.rs +++ b/crates/wasmtime/src/trampoline.rs @@ -41,6 +41,7 @@ fn create_handle( let handle = OnDemandInstanceAllocator::new(config.mem_creator.clone(), 0).allocate( InstanceAllocationRequest { module: Arc::new(module), + unique_id: None, memfds: None, functions, image_base: 0, diff --git a/crates/wasmtime/src/trampoline/func.rs b/crates/wasmtime/src/trampoline/func.rs index 47513f83cfea..77d5f26d188d 100644 --- a/crates/wasmtime/src/trampoline/func.rs +++ b/crates/wasmtime/src/trampoline/func.rs @@ -161,6 +161,7 @@ pub unsafe fn create_raw_function( Ok( OnDemandInstanceAllocator::default().allocate(InstanceAllocationRequest { module: Arc::new(module), + unique_id: None, memfds: None, functions: &functions, image_base: (*func).as_ptr() as usize,