T0b1/wasmtime
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Use deterministic randomness fuzzing the pooling allocator (#5247)

Browse Source 
This commit updates the index allocation performed in the pooling
allocator with a few refactorings:

* With `cfg(fuzzing)` a deterministic rng is now used to improve
  reproducibility of fuzz test cases.
* The `Mutex` was pushed inside of `IndexAllocator`, renamed from
  `PoolingAllocationState`.
* Randomness is now always done through a `SmallRng` stored in the
  `IndexAllocator` instead of using `thread_rng`.
* The `is_empty` method has been removed in favor of an `Option`-based
  return on `alloc`.

This refactoring is additionally intended to encapsulate more
implementation details of `IndexAllocator` to more easily allow for
alternate implementations in the future such as lock-free approaches
(possibly).
This commit is contained in:
Alex Crichton
2022-11-10 14:53:04 -06:00
committed by GitHub
parent 42e88c7b24
commit 7ec626b898
3 changed files with 141 additions and 125 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
crates/runtime/src/instance/allocator/pooling.rs
View File

@@ -24,7 +24,7 @@ use wasmtime_environ::{
};
mod index_allocator;
use index_allocator::{PoolingAllocationState, SlotId};
use index_allocator::{IndexAllocator, SlotId};
cfg_if::cfg_if! {
if #[cfg(windows)] {
@@ -119,7 +119,7 @@ struct InstancePool {
mapping: Mmap,
instance_size: usize,
max_instances: usize,
index_allocator: Mutex<PoolingAllocationState>,
index_allocator: IndexAllocator,
memories: MemoryPool,
tables: TablePool,
linear_memory_keep_resident: usize,
@@ -148,10 +148,7 @@ impl InstancePool {
mapping,
instance_size,
max_instances,
index_allocator: Mutex::new(PoolingAllocationState::new(
config.strategy,
max_instances,
)),
index_allocator: IndexAllocator::new(config.strategy, max_instances),
memories: MemoryPool::new(&config.limits, tunables)?,
tables: TablePool::new(&config.limits)?,
linear_memory_keep_resident: config.linear_memory_keep_resident,
@@ -221,21 +218,18 @@ impl InstancePool {
&self,
req: InstanceAllocationRequest,
) -> Result<InstanceHandle, InstantiationError> {
let index = {
let mut alloc = self.index_allocator.lock().unwrap();
if alloc.is_empty() {
return Err(InstantiationError::Limit(self.max_instances as u32));
}
alloc.alloc(req.runtime_info.unique_id()).index()
};
let id = self
.index_allocator
.alloc(req.runtime_info.unique_id())
.ok_or_else(|| InstantiationError::Limit(self.max_instances as u32))?;
match unsafe { self.initialize_instance(index, req) } {
match unsafe { self.initialize_instance(id.index(), req) } {
Ok(handle) => Ok(handle),
Err(e) => {
// If we failed to initialize the instance, there's no need to drop
// it as it was never "allocated", but we still need to free the
// instance's slot.
self.index_allocator.lock().unwrap().free(SlotId(index));
self.index_allocator.free(id);
Err(e)
}
}
@@ -267,7 +261,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(SlotId(index));
self.index_allocator.free(SlotId(index));
}
fn allocate_instance_resources(
@@ -839,7 +833,7 @@ struct StackPool {
stack_size: usize,
max_instances: usize,
page_size: usize,
index_allocator: Mutex<PoolingAllocationState>,
index_allocator: IndexAllocator,
async_stack_zeroing: bool,
async_stack_keep_resident: usize,
}
@@ -893,10 +887,10 @@ impl StackPool {
// 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(
index_allocator: IndexAllocator::new(
PoolingAllocationStrategy::NextAvailable,
max_instances,
)),
),
})
}
@@ -905,13 +899,11 @@ impl StackPool {
return Err(FiberStackError::NotSupported);
}
let index = {
let mut alloc = self.index_allocator.lock().unwrap();
if alloc.is_empty() {
return Err(FiberStackError::Limit(self.max_instances as u32));
}
alloc.alloc(None).index()
};
let index = self
.index_allocator
.alloc(None)
.ok_or(FiberStackError::Limit(self.max_instances as u32))?
.index();
assert!(index < self.max_instances);
@@ -958,7 +950,7 @@ impl StackPool {
self.zero_stack(bottom_of_stack, stack_size);
}
self.index_allocator.lock().unwrap().free(SlotId(index));
self.index_allocator.free(SlotId(index));
}
fn zero_stack(&self, bottom: usize, size: usize) {
@@ -1199,8 +1191,8 @@ mod test {
assert_eq!(instances.max_instances, 3);
assert_eq!(
instances.index_allocator.lock().unwrap().testing_freelist(),
&[SlotId(0), SlotId(1), SlotId(2)]
instances.index_allocator.testing_freelist(),
[SlotId(0), SlotId(1), SlotId(2)]
);
let mut handles = Vec::new();
@@ -1224,10 +1216,7 @@ mod test {
);
}
assert_eq!(
instances.index_allocator.lock().unwrap().testing_freelist(),
&[]
);
assert_eq!(instances.index_allocator.testing_freelist(), []);
match instances.allocate(InstanceAllocationRequest {
runtime_info: &empty_runtime_info(module),
@@ -1249,8 +1238,8 @@ mod test {
}
assert_eq!(
instances.index_allocator.lock().unwrap().testing_freelist(),
&[SlotId(2), SlotId(1), SlotId(0)]
instances.index_allocator.testing_freelist(),
[SlotId(2), SlotId(1), SlotId(0)]
);
Ok(())
@@ -1356,8 +1345,8 @@ mod test {
assert_eq!(pool.page_size, native_page_size);
assert_eq!(
pool.index_allocator.lock().unwrap().testing_freelist(),
&[
pool.index_allocator.testing_freelist(),
[
SlotId(0),
SlotId(1),
SlotId(2),
@@ -1383,7 +1372,7 @@ mod test {
stacks.push(stack);
}
assert_eq!(pool.index_allocator.lock().unwrap().testing_freelist(), &[]);
assert_eq!(pool.index_allocator.testing_freelist(), []);
match pool.allocate().unwrap_err() {
FiberStackError::Limit(10) => {}
@@ -1395,8 +1384,8 @@ mod test {
}
assert_eq!(
pool.index_allocator.lock().unwrap().testing_freelist(),
&[
pool.index_allocator.testing_freelist(),
[
SlotId(9),
SlotId(8),
SlotId(7),

195
crates/runtime/src/instance/allocator/pooling/index_allocator.rs
View File

@@ -2,8 +2,10 @@
use super::PoolingAllocationStrategy;
use crate::CompiledModuleId;
use rand::Rng;
use rand::rngs::SmallRng;
use rand::{Rng, SeedableRng};
use std::collections::HashMap;
use std::sync::Mutex;
/// A slot index. The job of this allocator is to hand out these
/// indices.
@@ -36,8 +38,17 @@ impl PerModuleFreeListIndex {
}
}
#[derive(Clone, Debug)]
pub(crate) enum PoolingAllocationState {
#[derive(Debug)]
pub struct IndexAllocator(Mutex<Inner>);
#[derive(Debug)]
struct Inner {
rng: SmallRng,
state: State,
}
#[derive(Debug)]
enum State {
NextAvailable(Vec<SlotId>),
Random(Vec<SlotId>),
/// Reuse-affinity policy state.
@@ -246,14 +257,14 @@ fn remove_module_free_list_item(
}
}
impl PoolingAllocationState {
impl IndexAllocator {
/// Create the default state for this strategy.
pub(crate) fn new(strategy: PoolingAllocationStrategy, max_instances: usize) -> Self {
pub 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 {
let state = match strategy {
PoolingAllocationStrategy::NextAvailable => State::NextAvailable(ids),
PoolingAllocationStrategy::Random => State::Random(ids),
PoolingAllocationStrategy::ReuseAffinity => State::ReuseAffinity {
free_list: ids,
per_module: HashMap::new(),
slot_state: (0..max_instances)
@@ -264,35 +275,37 @@ impl PoolingAllocationState {
})
.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(),
}
};
// Use a deterministic seed during fuzzing to improve reproducibility of
// test cases, but otherwise outside of fuzzing use a random seed to
// shake things up.
let seed = if cfg!(fuzzing) {
[0; 32]
} else {
rand::thread_rng().gen()
};
let rng = SmallRng::from_seed(seed);
IndexAllocator(Mutex::new(Inner { rng, state }))
}
/// 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()
pub fn alloc(&self, id: Option<CompiledModuleId>) -> Option<SlotId> {
let mut inner = self.0.lock().unwrap();
let inner = &mut *inner;
match &mut inner.state {
State::NextAvailable(free_list) => free_list.pop(),
State::Random(free_list) => {
if free_list.len() == 0 {
None
} else {
let id = inner.rng.gen_range(0..free_list.len());
Some(free_list.swap_remove(id))
}
}
&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,
..
State::ReuseAffinity {
free_list,
per_module,
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
@@ -308,8 +321,11 @@ impl PoolingAllocationState {
// 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
Some(slot_id)
} else {
if free_list.len() == 0 {
return None;
}
// Pick a random free slot ID. Note that we do
// this, rather than pick a victim module first,
// to maintain an unbiased stealing distribution:
@@ -333,7 +349,7 @@ impl PoolingAllocationState {
// 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 free_list_index = inner.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.
@@ -345,22 +361,22 @@ impl PoolingAllocationState {
}
slot_state[slot_id.index()] = SlotState::Taken(id);
slot_id
Some(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) => {
pub(crate) fn free(&self, index: SlotId) {
let mut inner = self.0.lock().unwrap();
match &mut inner.state {
State::NextAvailable(free_list) | State::Random(free_list) => {
free_list.push(index);
}
&mut PoolingAllocationState::ReuseAffinity {
ref mut per_module,
ref mut free_list,
ref mut slot_state,
State::ReuseAffinity {
per_module,
free_list,
slot_state,
} => {
let module_id = slot_state[index.index()].unwrap_module_id();
@@ -388,10 +404,10 @@ 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) -> &[SlotId] {
match self {
&PoolingAllocationState::NextAvailable(ref free_list)
| &PoolingAllocationState::Random(ref free_list) => &free_list[..],
pub(crate) fn testing_freelist(&self) -> Vec<SlotId> {
let inner = self.0.lock().unwrap();
match &inner.state {
State::NextAvailable(free_list) | State::Random(free_list) => free_list.clone(),
_ => panic!("Wrong kind of state"),
}
}
@@ -400,11 +416,12 @@ impl PoolingAllocationState {
/// 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(..) => {
let inner = self.0.lock().unwrap();
match &inner.state {
State::NextAvailable(..) | State::Random(..) => {
panic!("Wrong kind of state")
}
&PoolingAllocationState::ReuseAffinity { ref per_module, .. } => {
State::ReuseAffinity { per_module, .. } => {
let mut ret = vec![];
for (module, list) in per_module {
assert!(!list.is_empty());
@@ -418,28 +435,34 @@ impl PoolingAllocationState {
#[cfg(test)]
mod test {
use super::{PoolingAllocationState, SlotId};
use super::{IndexAllocator, 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);
for size in 0..20 {
let state = IndexAllocator::new(strat, size);
for i in 0..size {
assert_eq!(state.alloc(None).unwrap().index(), size - i - 1);
}
assert!(state.alloc(None).is_none());
}
}
#[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);
for size in 0..20 {
let state = IndexAllocator::new(strat, size);
for _ in 0..size {
assert!(state.alloc(None).unwrap().index() < size);
}
assert!(state.alloc(None).is_none());
}
}
#[test]
@@ -448,16 +471,16 @@ mod test {
let id_alloc = CompiledModuleIdAllocator::new();
let id1 = id_alloc.alloc();
let id2 = id_alloc.alloc();
let mut state = PoolingAllocationState::new(strat, 100);
let state = IndexAllocator::new(strat, 100);
let index1 = state.alloc(Some(id1));
let index1 = state.alloc(Some(id1)).unwrap();
assert!(index1.index() < 100);
let index2 = state.alloc(Some(id2));
let index2 = state.alloc(Some(id2)).unwrap();
assert!(index2.index() < 100);
assert_ne!(index1, index2);
state.free(index1);
let index3 = state.alloc(Some(id1));
let index3 = state.alloc(Some(id1)).unwrap();
assert_eq!(index3, index1);
state.free(index3);
@@ -476,10 +499,9 @@ mod test {
let mut indices = vec![];
for _ in 0..100 {
assert!(!state.is_empty());
indices.push(state.alloc(Some(id2)));
indices.push(state.alloc(Some(id2)).unwrap());
}
assert!(state.is_empty());
assert!(state.alloc(None).is_none());
assert_eq!(indices[0], index2);
for i in indices {
@@ -493,7 +515,7 @@ mod test {
// Allocate an index we know previously had an instance but
// now does not (list ran empty).
let index = state.alloc(Some(id1));
let index = state.alloc(Some(id1)).unwrap();
state.free(index);
}
@@ -507,25 +529,30 @@ mod test {
let ids = std::iter::repeat_with(|| id_alloc.alloc())
.take(10)
.collect::<Vec<_>>();
let mut state = PoolingAllocationState::new(strat, 1000);
let state = IndexAllocator::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;
loop {
if !allocated.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 {
let id = ids[rng.gen_range(0..ids.len())];
let index = match state.alloc(Some(id)) {
Some(id) => id,
None => continue,
};
if last_id[index.index()] == Some(id) {
hits += 1;
}
last_id[index.index()] = Some(id);
allocated.push(index);
}
last_id[index.index()] = Some(id);
allocated.push(index);
break;
}
}