Chris Fallin
3 years ago
committed by
GitHub
GitHub
11 changed files with 723 additions and 128 deletions
@ -0,0 +1,541 @@ |
|||
//! Index/slot allocator policies for the pooling allocator.
|
|||
|
|||
use super::PoolingAllocationStrategy; |
|||
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<SlotId>), |
|||
Random(Vec<SlotId>), |
|||
/// 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 (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.
|
|||
/// - 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<SlotId>, |
|||
/// 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<CompiledModuleId, Vec<SlotId>>, |
|||
/// 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<SlotState>, |
|||
}, |
|||
} |
|||
|
|||
#[derive(Clone, Debug)] |
|||
pub(crate) enum SlotState { |
|||
/// Currently allocated.
|
|||
///
|
|||
/// Invariant: no slot in this state has its index in either
|
|||
/// `free_list` or any list in `per_module`.
|
|||
Taken(Option<CompiledModuleId>), |
|||
/// 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<CompiledModuleId> { |
|||
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: 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: PerModuleFreeListIndex, |
|||
}, |
|||
} |
|||
|
|||
impl FreeSlotState { |
|||
/// Get the index of this slot in the global free list.
|
|||
fn free_list_index(&self) -> GlobalFreeListIndex { |
|||
match self { |
|||
&Self::NoAffinity { free_list_index } |
|||
| &Self::Affinity { |
|||
free_list_index, .. |
|||
} => free_list_index, |
|||
} |
|||
} |
|||
|
|||
/// Update the index of this slot in the global free list.
|
|||
fn update_free_list_index(&mut self, index: GlobalFreeListIndex) { |
|||
match self { |
|||
&mut Self::NoAffinity { |
|||
ref mut free_list_index, |
|||
} |
|||
| &mut Self::Affinity { |
|||
ref mut free_list_index, |
|||
.. |
|||
} => { |
|||
*free_list_index = index; |
|||
} |
|||
} |
|||
} |
|||
|
|||
/// Get the index of this slot in its per-module free list.
|
|||
fn per_module_index(&self) -> PerModuleFreeListIndex { |
|||
match self { |
|||
&Self::Affinity { |
|||
per_module_index, .. |
|||
} => per_module_index, |
|||
_ => 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: PerModuleFreeListIndex) { |
|||
match self { |
|||
&mut Self::Affinity { |
|||
ref mut per_module_index, |
|||
.. |
|||
} => { |
|||
*per_module_index = index; |
|||
} |
|||
_ => 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<SlotState>, |
|||
free_list: &mut Vec<SlotId>, |
|||
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<SlotState>, |
|||
per_module: &mut HashMap<CompiledModuleId, Vec<SlotId>>, |
|||
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).map(|i| SlotId(i)).collect::<Vec<_>>(); |
|||
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::Free(FreeSlotState::NoAffinity { |
|||
free_list_index: GlobalFreeListIndex(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(), |
|||
} |
|||
} |
|||
|
|||
/// Allocate a new slot.
|
|||
pub(crate) fn alloc(&mut self, id: Option<CompiledModuleId>) -> SlotId { |
|||
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 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.
|
|||
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,
|
|||
// 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 slot_id = free_list[free_list_index]; |
|||
// Remove from both the global freelist and
|
|||
// per-module freelist, if any.
|
|||
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[slot_id.index()] = SlotState::Taken(id); |
|||
|
|||
slot_id |
|||
} |
|||
} |
|||
} |
|||
} |
|||
|
|||
pub(crate) fn free(&mut self, index: SlotId) { |
|||
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 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) = 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.index()] = SlotState::Free(FreeSlotState::Affinity { |
|||
module: id, |
|||
free_list_index, |
|||
per_module_index, |
|||
}); |
|||
} else { |
|||
slot_state[index.index()] = |
|||
SlotState::Free(FreeSlotState::NoAffinity { 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) -> &[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<CompiledModuleId> { |
|||
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, SlotId}; |
|||
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).index(), 9); |
|||
let mut state = PoolingAllocationState::new(strat, 5); |
|||
assert_eq!(state.alloc(None).index(), 4); |
|||
let mut state = PoolingAllocationState::new(strat, 1); |
|||
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).index() < 100); |
|||
let mut state = PoolingAllocationState::new(strat, 1); |
|||
assert_eq!(state.alloc(None).index(), 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.index() < 100); |
|||
let index2 = state.alloc(Some(id2)); |
|||
assert!(index2.index() < 100); |
|||
assert_ne!(index1, index2); |
|||
|
|||
state.free(index1); |
|||
let index3 = state.alloc(Some(id1)); |
|||
assert_eq!(index3, index1); |
|||
state.free(index3); |
|||
|
|||
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
|
|||
// 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); |
|||
} |
|||
|
|||
// 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); |
|||
} |
|||
|
|||
#[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::<Vec<_>>(); |
|||
let mut state = PoolingAllocationState::new(strat, 1000); |
|||
let mut allocated: Vec<SlotId> = 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 = 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.index()] == Some(id) { |
|||
hits += 1; |
|||
} |
|||
last_id[index.index()] = Some(id); |
|||
allocated.push(index); |
|||
} |
|||
} |
|||
|
|||
// 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 |
|||
); |
|||
} |
|||
} |
|||
Loading…
Reference in new issue