Fixes #452: Introduce a random slot allocation strategy to Regions. (#496)

Author: fst-crenshaw

Cargo.lock

@@ -24,6 +24,7 @@ memoffset = "0.5.3"
 nix = "0.17"
 num-derive = "0.3.0"
 num-traits = "0.2"
+rand = "0.7"
 raw-cpuid = "6.0.0"
 thiserror = "1.0.4"
 tracing = "0.1.12"

lucet-runtime/lucet-runtime-internals/Cargo.toml

@@ -4,7 +4,8 @@ use crate::region::RegionInternal;
 use libc::c_void;
 use lucet_module::GlobalValue;
 use nix::unistd::{sysconf, SysconfVar};
-use std::sync::{Arc, Once, Weak};
+use rand::{thread_rng, Rng, RngCore};
+use std::sync::{Arc, Mutex, Once, Weak};
 
 pub const HOST_PAGE_SIZE_EXPECTED: usize = 4096;
 static mut HOST_PAGE_SIZE: usize = 0;
@@ -94,6 +95,45 @@ impl Slot {
     }
 }
 
+/// The strategy by which a `Region` selects an allocation to back an `Instance`.
+pub enum AllocStrategy {
+    /// Allocate from the next slot available.
+    Linear,
+    ///  Allocate randomly from the set of available slots.
+    Random,
+    /// Allocate randomly from the set of available slots using the
+    /// supplied random number generator.
+    ///
+    /// This strategy is used to create reproducible behavior for testing.
+    CustomRandom(Arc<Mutex<dyn RngCore>>),
+}
+
+impl AllocStrategy {
+    /// For a given `AllocStrategy`, use the number of free_slots and
+    /// capacity to determine the next slot to allocate for an
+    /// `Instance`.
+    pub fn next(&mut self, free_slots: usize, capacity: usize) -> Result<usize, Error> {
+        if free_slots == 0 {
+            return Err(Error::RegionFull(capacity));
+        }
+        match self {
+            AllocStrategy::Linear => Ok(free_slots - 1),
+            AllocStrategy::Random => {
+                // Instantiate a random number generator and get a
+                // random slot index.
+                let mut rng = thread_rng();
+                Ok(rng.gen_range(0, free_slots))
+            }
+            AllocStrategy::CustomRandom(custom_rng) => {
+                // Get a random slot index using the supplied random
+                // number generator.
+                let mut rng = custom_rng.lock().unwrap();
+                Ok(rng.gen_range(0, free_slots))
+            }
+        }
+    }
+}
+
 /// The structure that manages the allocations backing an `Instance`.
 ///
 /// `Alloc`s are not to be created directly, but rather are created by `Region`s during instance

lucet-runtime/lucet-runtime-internals/src/alloc/mod.rs

@@ -2,9 +2,11 @@
 macro_rules! alloc_tests {
     ( $TestRegion:path ) => {
         use libc::c_void;
-        use std::sync::Arc;
+        use rand::rngs::StdRng;
+        use rand::{thread_rng, Rng, SeedableRng};
+        use std::sync::{Arc, Mutex};
         use $TestRegion as TestRegion;
-        use $crate::alloc::{host_page_size, Limits, MINSIGSTKSZ};
+        use $crate::alloc::{host_page_size, AllocStrategy, Limits, MINSIGSTKSZ};
         use $crate::context::{Context, ContextHandle};
         use $crate::error::Error;
         use $crate::instance::InstanceInternal;
@@ -774,6 +776,173 @@ macro_rules! alloc_tests {
             assert_eq!(region.used_slots(), 0);
         }
 
+        /// This test exercises the AllocStrategy::Random. In this scenario,
+        /// the Region has a single slot which is "randomly" allocated and then dropped.
+        #[test]
+        fn slot_counts_work_with_random_alloc() {
+            let module = MockModuleBuilder::new()
+                .with_heap_spec(ONE_PAGE_HEAP)
+                .build();
+            let region = TestRegion::create(1, &LIMITS).expect("region created");
+            assert_eq!(region.capacity(), 1);
+            assert_eq!(region.free_slots(), 1);
+            assert_eq!(region.used_slots(), 0);
+
+            let inst = region
+                .new_instance_builder(module.clone())
+                .with_alloc_strategy(AllocStrategy::Random)
+                .build()
+                .expect("new_instance succeeds");
+            assert_eq!(region.capacity(), 1);
+            assert_eq!(region.free_slots(), 0);
+            assert_eq!(region.used_slots(), 1);
+
+            drop(inst);
+            assert_eq!(region.capacity(), 1);
+            assert_eq!(region.free_slots(), 1);
+            assert_eq!(region.used_slots(), 0);
+        }
+
+        /// This test exercises the AllocStrategy::CustomRandom. In this scenario,
+        /// the Region has 10 slots which are randomly allocated up to capacity
+        /// and then dropped. The test is executed 100 times to exercise the
+        /// random nature of the allocation strategy.
+        #[test]
+        fn slot_counts_work_with_custom_random_alloc() {
+            let mut master_rng = thread_rng();
+            let seed: u64 = master_rng.gen();
+            eprintln!(
+                "Seeding slot_counts_work_with_custom_random_alloc() with {}",
+                seed
+            );
+
+            let rng: StdRng = SeedableRng::seed_from_u64(seed);
+            let shared_rng = Arc::new(Mutex::new(rng));
+
+            for _ in 0..100 {
+                let mut inst_vec = Vec::new();
+                let module = MockModuleBuilder::new()
+                    .with_heap_spec(ONE_PAGE_HEAP)
+                    .build();
+                let total_slots = 10;
+                let region = TestRegion::create(total_slots, &LIMITS).expect("region created");
+                assert_eq!(region.capacity(), total_slots);
+                assert_eq!(region.free_slots(), 10);
+                assert_eq!(region.used_slots(), 0);
+
+                // Randomly allocate all of the slots in the region.
+                for i in 1..=total_slots {
+                    let inst = region
+                        .new_instance_builder(module.clone())
+                        .with_alloc_strategy(AllocStrategy::CustomRandom(shared_rng.clone()))
+                        .build()
+                        .expect("new_instance succeeds");
+
+                    assert_eq!(region.capacity(), total_slots);
+                    assert_eq!(region.free_slots(), total_slots - i);
+                    assert_eq!(region.used_slots(), i);
+                    inst_vec.push(inst);
+                }
+
+                // It's not possible to allocate just one more.  Try
+                // it and affirm that the error is handled gracefully.
+                let wont_inst = region
+                    .new_instance_builder(module.clone())
+                    .with_alloc_strategy(AllocStrategy::CustomRandom(shared_rng.clone()))
+                    .build();
+                assert!(wont_inst.is_err());
+
+                // Drop all of the slots in the region.
+                for i in 1..=total_slots {
+                    drop(inst_vec.pop());
+                    assert_eq!(region.capacity(), total_slots);
+                    assert_eq!(region.free_slots(), total_slots - (total_slots - i));
+                    assert_eq!(region.used_slots(), total_slots - i);
+                }
+
+                // Allocate just one more to make sure the drops took place
+                // and the Region has capacity again.
+                region
+                    .new_instance_builder(module.clone())
+                    .with_alloc_strategy(AllocStrategy::CustomRandom(shared_rng.clone()))
+                    .build()
+                    .expect("new_instance succeeds");
+            }
+        }
+
+        /// This test exercises a mixed AllocStrategy. In this scenario,
+        /// the Region has 10 slots which are randomly and linearly allocated
+        /// up to capacity and then dropped. The test is executed 100 times to
+        /// exercise the random nature of the allocation strategy.
+        #[test]
+        fn slot_counts_work_with_mixed_alloc() {
+            let mut master_rng = thread_rng();
+            let seed: u64 = master_rng.gen();
+            eprintln!("Seeding slot_counts_work_with_mixed_alloc() with {}", seed);
+
+            let rng: StdRng = SeedableRng::seed_from_u64(seed);
+            let shared_rng = Arc::new(Mutex::new(rng));
+
+            for _ in 0..100 {
+                let mut inst_vec = Vec::new();
+                let module = MockModuleBuilder::new()
+                    .with_heap_spec(ONE_PAGE_HEAP)
+                    .build();
+                let total_slots = 10;
+                let region = TestRegion::create(total_slots, &LIMITS).expect("region created");
+                assert_eq!(region.capacity(), total_slots);
+                assert_eq!(region.free_slots(), 10);
+                assert_eq!(region.used_slots(), 0);
+
+                // Allocate all of the slots in the region.
+                for i in 1..=total_slots {
+                    let inst;
+                    if i % 2 == 0 {
+                        inst = region
+                            .new_instance_builder(module.clone())
+                            .with_alloc_strategy(AllocStrategy::CustomRandom(shared_rng.clone()))
+                            .build()
+                            .expect("new_instance succeeds");
+                    } else {
+                        inst = region
+                            .new_instance_builder(module.clone())
+                            .with_alloc_strategy(AllocStrategy::Linear)
+                            .build()
+                            .expect("new_instance succeeds");
+                    }
+
+                    assert_eq!(region.capacity(), total_slots);
+                    assert_eq!(region.free_slots(), total_slots - i);
+                    assert_eq!(region.used_slots(), i);
+                    inst_vec.push(inst);
+                }
+
+                // It's not possible to allocate just one more.  Try
+                // it and affirm that the error is handled gracefully.
+                let wont_inst = region
+                    .new_instance_builder(module.clone())
+                    .with_alloc_strategy(AllocStrategy::CustomRandom(shared_rng.clone()))
+                    .build();
+                assert!(wont_inst.is_err());
+
+                // Drop all of the slots in the region.
+                for i in 1..=total_slots {
+                    drop(inst_vec.pop());
+                    assert_eq!(region.capacity(), total_slots);
+                    assert_eq!(region.free_slots(), total_slots - (total_slots - i));
+                    assert_eq!(region.used_slots(), total_slots - i);
+                }
+
+                // Allocate just one more to make sure the drops took place
+                // and the Region has capacity again.
+                region
+                    .new_instance_builder(module.clone())
+                    .with_alloc_strategy(AllocStrategy::Linear)
+                    .build()
+                    .expect("new_instance succeeds");
+            }
+        }
+
         fn do_nothing_module() -> Arc<dyn Module> {
             extern "C" fn do_nothing(_vmctx: *mut lucet_vmctx) -> () {}
 

lucet-runtime/lucet-runtime-internals/src/alloc/tests.rs

@@ -107,7 +107,18 @@ pub trait ModuleInternal: Send + Sync {
     ///
     /// Returns a `Result<(), Error>` rather than a boolean in order to provide a richer accounting
     /// of what may be invalid.
-    fn validate_runtime_spec(&self, limits: &Limits) -> Result<(), Error> {
+    fn validate_runtime_spec(
+        &self,
+        limits: &Limits,
+        instance_heap_limit: usize,
+    ) -> Result<(), Error> {
+        if instance_heap_limit > limits.heap_memory_size {
+            return Err(Error::InvalidArgument(
+                "heap memory size requested for instance is larger than slot allows",
+            ));
+        }
+        let heap_memory_size = std::cmp::min(limits.heap_memory_size, instance_heap_limit);
+
         // Modules without heap specs will not access the heap
         if let Some(heap) = self.heap_spec() {
             // Assure that the total reserved + guard regions fit in the address space.
@@ -128,7 +139,7 @@ pub trait ModuleInternal: Send + Sync {
                 bail_limits_exceeded!("heap spec reserved and guard size: {:?}", heap);
             }
 
-            if heap.initial_size as usize > limits.heap_memory_size {
+            if heap.initial_size as usize > heap_memory_size {
                 bail_limits_exceeded!("heap spec initial size: {:?}", heap);
             }
 

lucet-runtime/lucet-runtime-internals/src/module.rs

@@ -1,6 +1,6 @@
 pub mod mmap;
 
-use crate::alloc::{Alloc, Limits, Slot};
+use crate::alloc::{Alloc, AllocStrategy, Limits, Slot};
 use crate::embed_ctx::CtxMap;
 use crate::error::Error;
 use crate::instance::InstanceHandle;
@@ -53,6 +53,7 @@ pub trait RegionInternal: Send + Sync {
         module: Arc<dyn Module>,
         embed_ctx: CtxMap,
         heap_memory_size_limit: usize,
+        alloc_strategy: AllocStrategy,
     ) -> Result<InstanceHandle, Error>;
 
     /// Unmaps the heap, stack, and globals of an `Alloc`, while retaining the virtual address
@@ -90,6 +91,7 @@ pub struct InstanceBuilder<'a> {
     module: Arc<dyn Module>,
     embed_ctx: CtxMap,
     heap_memory_size_limit: usize,
+    alloc_strategy: AllocStrategy,
 }
 
 impl<'a> InstanceBuilder<'a> {
@@ -99,9 +101,21 @@ impl<'a> InstanceBuilder<'a> {
             module,
             embed_ctx: CtxMap::default(),
             heap_memory_size_limit: region.get_limits().heap_memory_size,
+            alloc_strategy: AllocStrategy::Linear,
         }
     }
 
+    /// Allocate the instance using the supplied `AllocStrategy`.
+    ///
+    /// This call is optional.  The default allocation strategy for
+    /// Regions is Linear, which allocates the instance using next available
+    /// alloc.  If a different strategy is desired, choose from those
+    /// available in `AllocStrategy`.
+    pub fn with_alloc_strategy(mut self, alloc_strategy: AllocStrategy) -> Self {
+        self.alloc_strategy = alloc_strategy;
+        self
+    }
+
     /// Add a smaller, custom limit for the heap memory size to the built instance.
     ///
     /// This call is optional. Attempts to build a new instance fail if the
@@ -122,7 +136,11 @@ impl<'a> InstanceBuilder<'a> {
 
     /// Build the instance.
     pub fn build(self) -> Result<InstanceHandle, Error> {
-        self.region
-            .new_instance_with(self.module, self.embed_ctx, self.heap_memory_size_limit)
+        self.region.new_instance_with(
+            self.module,
+            self.embed_ctx,
+            self.heap_memory_size_limit,
+            self.alloc_strategy,
+        )
     }
 }