From 1cbd39393074e71c06f4aec8ed305a814d78ac85 Mon Sep 17 00:00:00 2001 From: Chris Fallin Date: Mon, 31 Jan 2022 14:39:21 -0800 Subject: [PATCH] Review comments. --- .../runtime/src/instance/allocator/pooling.rs | 41 +- .../allocator/pooling/index_allocator.rs | 361 ++++++++++++------ 2 files changed, 278 insertions(+), 124 deletions(-) diff --git a/crates/runtime/src/instance/allocator/pooling.rs b/crates/runtime/src/instance/allocator/pooling.rs index b5685c876bcf..90db21fdebcb 100644 --- a/crates/runtime/src/instance/allocator/pooling.rs +++ b/crates/runtime/src/instance/allocator/pooling.rs @@ -25,7 +25,7 @@ use wasmtime_environ::{ }; mod index_allocator; -use index_allocator::PoolingAllocationState; +use index_allocator::{PoolingAllocationState, SlotId}; cfg_if::cfg_if! { if #[cfg(windows)] { @@ -404,7 +404,7 @@ impl InstancePool { if alloc.is_empty() { return Err(InstantiationError::Limit(self.max_instances as u32)); } - alloc.alloc(req.unique_id) + alloc.alloc(req.unique_id).index() }; unsafe { @@ -430,7 +430,6 @@ 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 @@ -500,7 +499,7 @@ impl InstancePool { // touched again until we write a fresh Instance in-place with // std::ptr::write in allocate() above. - self.index_allocator.lock().unwrap().free(index, unique_id); + self.index_allocator.lock().unwrap().free(SlotId(index)); } fn set_instance_memories( @@ -928,7 +927,7 @@ impl StackPool { if alloc.is_empty() { return Err(FiberStackError::Limit(self.max_instances as u32)); } - alloc.alloc(None) + alloc.alloc(None).index() }; debug_assert!(index < self.max_instances); @@ -974,7 +973,7 @@ impl StackPool { decommit_stack_pages(bottom_of_stack as _, stack_size).unwrap(); - self.index_allocator.lock().unwrap().free(index, None); + self.index_allocator.lock().unwrap().free(SlotId(index)); } } @@ -1457,7 +1456,7 @@ mod test { assert_eq!( instances.index_allocator.lock().unwrap().testing_freelist(), - &[0, 1, 2] + &[SlotId(0), SlotId(1), SlotId(2)] ); let mut handles = Vec::new(); @@ -1520,7 +1519,7 @@ mod test { assert_eq!( instances.index_allocator.lock().unwrap().testing_freelist(), - &[2, 1, 0] + &[SlotId(2), SlotId(1), SlotId(0)] ); Ok(()) @@ -1632,7 +1631,18 @@ mod test { assert_eq!( pool.index_allocator.lock().unwrap().testing_freelist(), - &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], + &[ + SlotId(0), + SlotId(1), + SlotId(2), + SlotId(3), + SlotId(4), + SlotId(5), + SlotId(6), + SlotId(7), + SlotId(8), + SlotId(9) + ], ); let base = pool.mapping.as_ptr() as usize; @@ -1660,7 +1670,18 @@ mod test { assert_eq!( pool.index_allocator.lock().unwrap().testing_freelist(), - &[9, 8, 7, 6, 5, 4, 3, 2, 1, 0], + &[ + SlotId(9), + SlotId(8), + SlotId(7), + SlotId(6), + SlotId(5), + SlotId(4), + SlotId(3), + SlotId(2), + SlotId(1), + SlotId(0) + ], ); Ok(()) diff --git a/crates/runtime/src/instance/allocator/pooling/index_allocator.rs b/crates/runtime/src/instance/allocator/pooling/index_allocator.rs index ee4485b8edb9..e2b7f13e93ee 100644 --- a/crates/runtime/src/instance/allocator/pooling/index_allocator.rs +++ b/crates/runtime/src/instance/allocator/pooling/index_allocator.rs @@ -5,16 +5,47 @@ use crate::CompiledModuleId; use rand::Rng; use std::collections::HashMap; +/// A slot index. The job of this allocator is to hand out these +/// indices. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct SlotId(pub usize); +impl SlotId { + /// The index of this slot. + pub fn index(self) -> usize { + self.0 + } +} + +/// An index in the global freelist. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct GlobalFreeListIndex(usize); +impl GlobalFreeListIndex { + /// The index of this slot. + fn index(self) -> usize { + self.0 + } +} + +/// An index in a per-module freelist. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct PerModuleFreeListIndex(usize); +impl PerModuleFreeListIndex { + /// The index of this slot. + fn index(self) -> usize { + self.0 + } +} + #[derive(Clone, Debug)] pub(crate) enum PoolingAllocationState { - NextAvailable(Vec), - Random(Vec), + 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. + /// matter their affinities (or no affinity at all). /// - 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. @@ -30,56 +61,112 @@ pub(crate) enum PoolingAllocationState { /// 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. + /// 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). We + /// remove a module's freelist when we have no more slots with + /// affinity for that module. + 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, + /// Currently allocated. + /// + /// Invariant: no slot in this state has its index in either + /// `free_list` or any list in `per_module`. + Taken(Option), + /// Currently free. A free slot is able to be allocated for any + /// request, but may have affinity to a certain module that we + /// prefer to use it for. + /// + /// Invariant: every slot in this state has its index in at least + /// `free_list`, and possibly a `per_module` free-list; see + /// FreeSlotState. + Free(FreeSlotState), +} + +impl SlotState { + fn unwrap_free(&self) -> &FreeSlotState { + match self { + &Self::Free(ref free) => free, + _ => panic!("Slot not free"), + } + } + + fn unwrap_free_mut(&mut self) -> &mut FreeSlotState { + match self { + &mut Self::Free(ref mut free) => free, + _ => panic!("Slot not free"), + } + } + + fn unwrap_module_id(&self) -> Option { + match self { + &Self::Taken(module_id) => module_id, + _ => panic!("Slot not in Taken state"), + } + } +} + +#[derive(Clone, Debug)] +pub(crate) enum FreeSlotState { + /// The slot is free, and has no affinity. + /// + /// Invariant: every slot in this state has its index in + /// `free_list`. No slot in this state has its index in any other + /// (per-module) free-list. + NoAffinity { + /// Index in the global free list. + /// + /// Invariant: free_list[slot_state[i].free_list_index] == i. + free_list_index: GlobalFreeListIndex, }, + /// The slot is free, and has an affinity for some module. This + /// means we prefer to choose this slot (or some other one with + /// the same affinity) given a request to allocate a slot for this + /// module. It can, however, still be used for any other module if + /// needed. + /// + /// Invariant: every slot in this state has its index in both + /// `free_list` *and* exactly one list in `per_module`. 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] + /// Index in the global free list. + /// + /// Invariant: free_list[slot_state[i].free_list_index] == i. + free_list_index: GlobalFreeListIndex, + /// 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, + per_module_index: PerModuleFreeListIndex, }, } -impl SlotState { +impl FreeSlotState { /// Get the index of this slot in the global free list. - fn free_list_index(&self) -> usize { + fn free_list_index(&self) -> GlobalFreeListIndex { match self { - &Self::Empty { free_list_index } + &Self::NoAffinity { 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) { + fn update_free_list_index(&mut self, index: GlobalFreeListIndex) { match self { - &mut Self::Empty { + &mut Self::NoAffinity { ref mut free_list_index, } | &mut Self::Affinity { @@ -88,22 +175,21 @@ impl SlotState { } => { *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 { + fn per_module_index(&self) -> PerModuleFreeListIndex { match self { &Self::Affinity { per_module_index, .. } => per_module_index, - _ => unreachable!(), + _ => panic!("per_module_index on slot with no affinity"), } } /// Update the index of this slot in its per-module free list. - fn update_per_module_index(&mut self, index: usize) { + fn update_per_module_index(&mut self, index: PerModuleFreeListIndex) { match self { &mut Self::Affinity { ref mut per_module_index, @@ -111,15 +197,59 @@ impl SlotState { } => { *per_module_index = index; } - _ => panic!("Taken in per-module free list"), + _ => panic!("per_module_index on slot with no affinity"), } } } +/// Internal: remove a slot-index from the global free list. +fn remove_global_free_list_item( + slot_state: &mut Vec, + free_list: &mut Vec, + index: SlotId, +) { + let free_list_index = slot_state[index.index()].unwrap_free().free_list_index(); + assert_eq!(index, free_list.swap_remove(free_list_index.index())); + if free_list_index.index() < free_list.len() { + let replaced = free_list[free_list_index.index()]; + slot_state[replaced.index()] + .unwrap_free_mut() + .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: SlotId, +) { + debug_assert!( + per_module.contains_key(&id), + "per_module list for given module should not be empty" + ); + + let per_module_list = per_module.get_mut(&id).unwrap(); + debug_assert!(!per_module_list.is_empty()); + + let per_module_index = slot_state[index.index()].unwrap_free().per_module_index(); + assert_eq!(index, per_module_list.swap_remove(per_module_index.index())); + if per_module_index.index() < per_module_list.len() { + let replaced = per_module_list[per_module_index.index()]; + slot_state[replaced.index()] + .unwrap_free_mut() + .update_per_module_index(per_module_index); + } + if per_module_list.is_empty() { + per_module.remove(&id); + } +} + 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::>(); + let ids = (0..max_instances).map(|i| SlotId(i)).collect::>(); match strategy { PoolingAllocationStrategy::NextAvailable => PoolingAllocationState::NextAvailable(ids), PoolingAllocationStrategy::Random => PoolingAllocationState::Random(ids), @@ -127,7 +257,11 @@ impl PoolingAllocationState { free_list: ids, per_module: HashMap::new(), slot_state: (0..max_instances) - .map(|i| SlotState::Empty { free_list_index: i }) + .map(|i| { + SlotState::Free(FreeSlotState::NoAffinity { + free_list_index: GlobalFreeListIndex(i), + }) + }) .collect(), }, } @@ -142,41 +276,8 @@ impl PoolingAllocationState { } } - /// 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 { + pub(crate) fn alloc(&mut self, id: Option) -> SlotId { match self { &mut PoolingAllocationState::NextAvailable(ref mut free_list) => { debug_assert!(free_list.len() > 0); @@ -198,16 +299,16 @@ impl PoolingAllocationState { // 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"); + let slot_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 + remove_global_free_list_item(slot_state, free_list, slot_id); + slot_state[slot_id.index()] = SlotState::Taken(id); + slot_id } else { // Pick a random free slot ID. Note that we do // this, rather than pick a victim module first, @@ -233,22 +334,24 @@ impl PoolingAllocationState { // (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]; + let slot_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); + remove_global_free_list_item(slot_state, free_list, slot_id); + if let &SlotState::Free(FreeSlotState::Affinity { module, .. }) = + &slot_state[slot_id.index()] + { + remove_module_free_list_item(slot_state, per_module, module, slot_id); } - slot_state[new_id] = SlotState::Taken; + slot_state[slot_id.index()] = SlotState::Taken(id); - new_id + slot_id } } } } - pub(crate) fn free(&mut self, index: usize, id: Option) { + pub(crate) fn free(&mut self, index: SlotId) { match self { &mut PoolingAllocationState::NextAvailable(ref mut free_list) | &mut PoolingAllocationState::Random(ref mut free_list) => { @@ -259,19 +362,24 @@ impl PoolingAllocationState { ref mut free_list, ref mut slot_state, } => { - let free_list_index = free_list.len(); + let module_id = slot_state[index.index()].unwrap_module_id(); + + let free_list_index = GlobalFreeListIndex(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(); + if let Some(id) = module_id { + let per_module_list = per_module + .entry(id) + .or_insert_with(|| Vec::with_capacity(1)); + let per_module_index = PerModuleFreeListIndex(per_module_list.len()); per_module_list.push(index); - slot_state[index] = SlotState::Affinity { + slot_state[index.index()] = SlotState::Free(FreeSlotState::Affinity { module: id, free_list_index, per_module_index, - }; + }); } else { - slot_state[index] = SlotState::Empty { free_list_index }; + slot_state[index.index()] = + SlotState::Free(FreeSlotState::NoAffinity { free_list_index }); } } } @@ -280,18 +388,37 @@ impl PoolingAllocationState { /// 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] { + pub(crate) fn testing_freelist(&self) -> &[SlotId] { match self { &PoolingAllocationState::NextAvailable(ref free_list) | &PoolingAllocationState::Random(ref free_list) => &free_list[..], _ => panic!("Wrong kind of state"), } } + + /// 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 { + match self { + &PoolingAllocationState::NextAvailable(..) | &PoolingAllocationState::Random(..) => { + panic!("Wrong kind of state") + } + &PoolingAllocationState::ReuseAffinity { ref per_module, .. } => { + let mut ret = vec![]; + for (module, list) in per_module { + assert!(!list.is_empty()); + ret.push(*module); + } + ret + } + } + } } #[cfg(test)] mod test { - use super::PoolingAllocationState; + use super::{PoolingAllocationState, SlotId}; use crate::CompiledModuleIdAllocator; use crate::PoolingAllocationStrategy; @@ -299,20 +426,20 @@ mod test { fn test_next_available_allocation_strategy() { let strat = PoolingAllocationStrategy::NextAvailable; let mut state = PoolingAllocationState::new(strat, 10); - assert_eq!(state.alloc(None), 9); + assert_eq!(state.alloc(None).index(), 9); let mut state = PoolingAllocationState::new(strat, 5); - assert_eq!(state.alloc(None), 4); + assert_eq!(state.alloc(None).index(), 4); let mut state = PoolingAllocationState::new(strat, 1); - assert_eq!(state.alloc(None), 0); + assert_eq!(state.alloc(None).index(), 0); } #[test] fn test_random_allocation_strategy() { let strat = PoolingAllocationStrategy::Random; let mut state = PoolingAllocationState::new(strat, 100); - assert!(state.alloc(None) < 100); + assert!(state.alloc(None).index() < 100); let mut state = PoolingAllocationState::new(strat, 1); - assert_eq!(state.alloc(None), 0); + assert_eq!(state.alloc(None).index(), 0); } #[test] @@ -324,17 +451,23 @@ mod test { let mut state = PoolingAllocationState::new(strat, 100); let index1 = state.alloc(Some(id1)); - assert!(index1 < 100); + assert!(index1.index() < 100); let index2 = state.alloc(Some(id2)); - assert!(index2 < 100); + assert!(index2.index() < 100); assert_ne!(index1, index2); - state.free(index1, Some(id1)); + state.free(index1); let index3 = state.alloc(Some(id1)); assert_eq!(index3, index1); - state.free(index3, Some(id1)); + state.free(index3); - state.free(index2, Some(id2)); + state.free(index2); + + // Both id1 and id2 should have some slots with affinity. + let affinity_modules = state.testing_module_affinity_list(); + assert_eq!(2, affinity_modules.len()); + assert!(affinity_modules.contains(&id1)); + assert!(affinity_modules.contains(&id2)); // Now there is 1 free instance for id2 and 1 free instance // for id1, and 98 empty. Allocate 100 for id2. The first @@ -350,13 +483,18 @@ mod test { assert_eq!(indices[0], index2); for i in indices { - state.free(i, Some(id2)); + state.free(i); } + // Now there should be no slots left with affinity for id1. + let affinity_modules = state.testing_module_affinity_list(); + assert_eq!(1, affinity_modules.len()); + assert!(affinity_modules.contains(&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)); + state.free(index); } #[test] @@ -370,29 +508,24 @@ mod test { .take(10) .collect::>(); let mut state = PoolingAllocationState::new(strat, 1000); - let mut allocated = vec![]; + let mut allocated: Vec = 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); + let to_free_idx = allocated.swap_remove(i); + state.free(to_free_idx); } 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) { + if last_id[index.index()] == Some(id) { hits += 1; } - last_id[index] = Some(id); - allocated.push((index, id)); + last_id[index.index()] = Some(id); + allocated.push(index); } }