From 01c044ee7d5c6b33431bc77499dfda808756a310 Mon Sep 17 00:00:00 2001 From: Chris Fallin Date: Fri, 28 Jan 2022 16:03:05 -0800 Subject: [PATCH] Pooling allocator: add a reuse-affinity policy. This policy attempts to reuse the same instance slot for subsequent instantiations of the same module. This is particularly useful when using a pooling backend such as memfd that benefits from this reuse: for example, in the memfd case, instantiating the same module into the same slot allows us to avoid several calls to mmap() because the same mappings can be reused. The policy tracks a freelist per "compiled module ID", and when allocating a slot for an instance, tries these two options in order: 1. A slot from the freelist for this module (i.e., last used for another instantiation of this particular module), or 2. A slot that was last used by some other module or never before. The "victim" slot for choice 2 is randomly chosen. The data structures are carefully designed so that all updates are O(1), and there is no retry-loop in any of the random selection. This policy is now the default when the memfd backend is selected via the `memfd-allocator` feature flag. --- crates/runtime/src/instance.rs | 7 +- crates/runtime/src/instance/allocator.rs | 15 +- .../runtime/src/instance/allocator/pooling.rs | 189 ++++---- .../allocator/pooling/index_allocator.rs | 408 ++++++++++++++++++ .../src/instance/allocator/pooling/uffd.rs | 43 +- crates/runtime/src/module_id.rs | 21 +- crates/wasmtime/src/config/pooling.rs | 6 + crates/wasmtime/src/instance.rs | 1 + crates/wasmtime/src/store.rs | 1 + crates/wasmtime/src/trampoline.rs | 1 + crates/wasmtime/src/trampoline/func.rs | 1 + 11 files changed, 567 insertions(+), 126 deletions(-) create mode 100644 crates/runtime/src/instance/allocator/pooling/index_allocator.rs 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,