Use deterministic randomness fuzzing the pooling allocator (#5247)

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).

crates/runtime/Cargo.toml

@@ -21,7 +21,7 @@ memoffset = "0.6.0"
 indexmap = "1.0.2"
 thiserror = "1.0.4"
 cfg-if = "1.0"
-rand = "0.8.3"
+rand = { version = "0.8.3", features = ['small_rng'] }
 anyhow = { workspace = true }
 memfd = "0.6.1"
 paste = "1.0.3"

crates/runtime/src/instance/allocator/pooling.rs

@@ -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(
+            index_allocator: IndexAllocator::new(config.strategy, max_instances),
-                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 id = self
-            let mut alloc = self.index_allocator.lock().unwrap();
+            .index_allocator
-            if alloc.is_empty() {
+            .alloc(req.runtime_info.unique_id())
-                return Err(InstantiationError::Limit(self.max_instances as u32));
+            .ok_or_else(|| InstantiationError::Limit(self.max_instances as u32))?;
-            }
-            alloc.alloc(req.runtime_info.unique_id()).index()
-        };
 
-        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 index = self
-            let mut alloc = self.index_allocator.lock().unwrap();
+            .index_allocator
-            if alloc.is_empty() {
+            .alloc(None)
-                return Err(FiberStackError::Limit(self.max_instances as u32));
+            .ok_or(FiberStackError::Limit(self.max_instances as u32))?
-            }
+            .index();
-            alloc.alloc(None).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(),
+            instances.index_allocator.testing_freelist(),
-            &[SlotId(0), SlotId(1), SlotId(2)]
+            [SlotId(0), SlotId(1), SlotId(2)]
         );
 
         let mut handles = Vec::new();
@@ -1224,10 +1216,7 @@ mod test {
             );
         }
 
-        assert_eq!(
+        assert_eq!(instances.index_allocator.testing_freelist(), []);
-            instances.index_allocator.lock().unwrap().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(),
+            instances.index_allocator.testing_freelist(),
-            &[SlotId(2), SlotId(1), SlotId(0)]
+            [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),

crates/runtime/src/instance/allocator/pooling/index_allocator.rs

@@ -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)]
+#[derive(Debug)]
-pub(crate) enum PoolingAllocationState {
+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 {
+        let state = match strategy {
-            PoolingAllocationStrategy::NextAvailable => PoolingAllocationState::NextAvailable(ids),
+            PoolingAllocationStrategy::NextAvailable => State::NextAvailable(ids),
-            PoolingAllocationStrategy::Random => PoolingAllocationState::Random(ids),
+            PoolingAllocationStrategy::Random => State::Random(ids),
-            PoolingAllocationStrategy::ReuseAffinity => PoolingAllocationState::ReuseAffinity {
+            PoolingAllocationStrategy::ReuseAffinity => State::ReuseAffinity {
                 free_list: ids,
                 per_module: HashMap::new(),
                 slot_state: (0..max_instances)
@@ -264,35 +275,37 @@ impl PoolingAllocationState {
                     })
                     .collect(),
             },
-        }
+        };
-    }
+        // Use a deterministic seed during fuzzing to improve reproducibility of
-
+        // test cases, but otherwise outside of fuzzing use a random seed to
-    /// Are any slots left, or is this allocator empty?
+        // shake things up.
-    pub(crate) fn is_empty(&self) -> bool {
+        let seed = if cfg!(fuzzing) {
-        match self {
+            [0; 32]
-            &PoolingAllocationState::NextAvailable(ref free_list)
+        } else {
-            | &PoolingAllocationState::Random(ref free_list) => free_list.is_empty(),
+            rand::thread_rng().gen()
-            &PoolingAllocationState::ReuseAffinity { ref free_list, .. } => free_list.is_empty(),
+        };
-        }
+        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 {
+    pub fn alloc(&self, id: Option<CompiledModuleId>) -> Option<SlotId> {
-        match self {
+        let mut inner = self.0.lock().unwrap();
-            &mut PoolingAllocationState::NextAvailable(ref mut free_list) => {
+        let inner = &mut *inner;
-                debug_assert!(free_list.len() > 0);
+        match &mut inner.state {
-                free_list.pop().unwrap()
+            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 {
+            State::ReuseAffinity {
-                ref mut free_list,
+                free_list,
-                ref mut per_module,
+                per_module,
-                ref mut slot_state,
+                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) {
+    pub(crate) fn free(&self, index: SlotId) {
-        match self {
+        let mut inner = self.0.lock().unwrap();
-            &mut PoolingAllocationState::NextAvailable(ref mut free_list)
+        match &mut inner.state {
-            | &mut PoolingAllocationState::Random(ref mut free_list) => {
+            State::NextAvailable(free_list) | State::Random(free_list) => {
                 free_list.push(index);
             }
-            &mut PoolingAllocationState::ReuseAffinity {
+            State::ReuseAffinity {
-                ref mut per_module,
+                per_module,
-                ref mut free_list,
+                free_list,
-                ref mut slot_state,
+                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] {
+    pub(crate) fn testing_freelist(&self) -> Vec<SlotId> {
-        match self {
+        let inner = self.0.lock().unwrap();
-            &PoolingAllocationState::NextAvailable(ref free_list)
+        match &inner.state {
-            | &PoolingAllocationState::Random(ref free_list) => &free_list[..],
+            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 {
+        let inner = self.0.lock().unwrap();
-            &PoolingAllocationState::NextAvailable(..) | &PoolingAllocationState::Random(..) => {
+        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);
+        for size in 0..20 {
-        let mut state = PoolingAllocationState::new(strat, 5);
+            let state = IndexAllocator::new(strat, size);
-        assert_eq!(state.alloc(None).index(), 4);
+            for i in 0..size {
-        let mut state = PoolingAllocationState::new(strat, 1);
+                assert_eq!(state.alloc(None).unwrap().index(), size - i - 1);
-        assert_eq!(state.alloc(None).index(), 0);
+            }
+            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);
+        for size in 0..20 {
-        let mut state = PoolingAllocationState::new(strat, 1);
+            let state = IndexAllocator::new(strat, size);
-        assert_eq!(state.alloc(None).index(), 0);
+            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)).unwrap());
-            indices.push(state.alloc(Some(id2)));
         }
-        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,26 +529,31 @@ 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)) {
+            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 {
-                assert!(!state.is_empty());
                     let id = ids[rng.gen_range(0..ids.len())];
-                let index = state.alloc(Some(id));
+                    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);
                 }
+                break;
+            }
         }
 
         // 10% reuse would be random chance (because we have 10 module