Merge pull request #77 from fastly/acf/lucet-wasi-memory-cli

Add CLI configuration for memory limits and improve performance for large initial heaps

Author: acfoltzer

Cargo.lock

@@ -388,6 +388,11 @@ impl Limits {
                 "address space size must be a multiple of host page size",
             ));
         }
+        if self.heap_memory_size > self.heap_address_space_size {
+            return Err(Error::InvalidArgument(
+                "address space size must be at least as large as memory size",
+            ));
+        }
         if self.stack_size % host_page_size() != 0 {
             return Err(Error::InvalidArgument(
                 "stack size must be a multiple of host page size",

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

@@ -256,6 +256,19 @@ macro_rules! alloc_tests {
             assert!(res.is_err(), "new_instance fails");
         }
 
+        /// This test shows that we reject limits with a larger memory size than address space size
+        #[test]
+        fn reject_undersized_address_space() {
+            const LIMITS: Limits = Limits {
+                heap_memory_size: LIMITS_HEAP_ADDRSPACE_SIZE + 4096,
+                heap_address_space_size: LIMITS_HEAP_ADDRSPACE_SIZE,
+                stack_size: LIMITS_STACK_SIZE,
+                globals_size: LIMITS_GLOBALS_SIZE,
+            };
+            let res = TestRegion::create(10, &LIMITS);
+            assert!(res.is_err(), "region creation fails");
+        }
+
         const SMALL_GUARD_HEAP: HeapSpec = HeapSpec {
             reserved_size: SPEC_HEAP_RESERVED_SIZE,
             guard_size: SPEC_HEAP_GUARD_SIZE - 1,

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

@@ -164,6 +164,13 @@ pub trait ModuleInternal: Send + Sync {
             bail_limits_exceeded!("heap spec initial size: {:?}", heap);
         }
 
+        if heap.initial_size > heap.reserved_size {
+            return Err(lucet_incorrect_module!(
+                "initial heap size greater than reserved size: {:?}",
+                heap
+            ));
+        }
+
         if self.globals().len() * std::mem::size_of::<u64>() > limits.globals_size {
             bail_limits_exceeded!("globals exceed limits");
         }

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

@@ -39,8 +39,6 @@ impl RegionInternal for MmapRegion {
         module.validate_runtime_spec(limits)?;
 
         for (ptr, len) in [
-            // make the heap read/writable and record its initial size
-            (slot.heap, module.heap_spec().initial_size as usize),
             // make the stack read/writable
             (slot.stack, limits.stack_size),
             // make the globals read/writable
@@ -54,6 +52,8 @@ impl RegionInternal for MmapRegion {
             unsafe { mprotect(*ptr, *len, ProtFlags::PROT_READ | ProtFlags::PROT_WRITE)? };
         }
 
+        // note: the initial heap will be made read/writable when `new_instance_handle` calls `reset`
+
         let inst_ptr = slot.start as *mut Instance;
 
         // upgrade the slot's weak region pointer so the region can't get dropped while the instance
@@ -66,7 +66,7 @@ impl RegionInternal for MmapRegion {
             .expect("backing region of slot (`self`) exists");
 
         let alloc = Alloc {
-            heap_accessible_size: module.heap_spec().initial_size as usize,
+            heap_accessible_size: 0, // the `reset` call in `new_instance_handle` will set this
             heap_inaccessible_size: slot.limits.heap_address_space_size,
             slot: Some(slot),
             region,
@@ -119,29 +119,32 @@ impl RegionInternal for MmapRegion {
     }
 
     fn reset_heap(&self, alloc: &mut Alloc, module: &dyn Module) -> Result<(), Error> {
+        let heap = alloc.slot().heap;
+
+        if alloc.heap_accessible_size > 0 {
+            // zero the whole heap, if any of it is currently accessible
+            let heap_size = alloc.slot().limits.heap_address_space_size;
+
+            unsafe {
+                mprotect(heap, heap_size, ProtFlags::PROT_NONE)?;
+                madvise(heap, heap_size, MmapAdvise::MADV_DONTNEED)?;
+            }
+        }
+
         let initial_size = module.heap_spec().initial_size as usize;
 
-        // reset the heap to the initial size
+        // reset the heap to the initial size, and mprotect those pages appropriately
         if alloc.heap_accessible_size != initial_size {
+            unsafe {
+                mprotect(
+                    heap,
+                    initial_size,
+                    ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
+                )?
+            };
             alloc.heap_accessible_size = initial_size;
             alloc.heap_inaccessible_size =
                 alloc.slot().limits.heap_address_space_size - initial_size;
-
-            // turn off any extra pages
-            let acc_heap_end =
-                (alloc.slot().heap as usize + alloc.heap_accessible_size) as *mut c_void;
-            unsafe {
-                mprotect(
-                    acc_heap_end,
-                    alloc.heap_inaccessible_size,
-                    ProtFlags::PROT_NONE,
-                )?;
-                madvise(
-                    acc_heap_end,
-                    alloc.heap_inaccessible_size,
-                    MmapAdvise::MADV_DONTNEED,
-                )?;
-            }
         }
 
         // Initialize the heap using the module sparse page data. There cannot be more pages in the
@@ -171,12 +174,6 @@ impl RegionInternal for MmapRegion {
             if let Some(contents) = module.get_sparse_page_data(page_num) {
                 // otherwise copy in the page data
                 heap[page_base..page_base + host_page_size()].copy_from_slice(contents);
-            } else {
-                // zero this page
-                // TODO would using a memset primitive be better here?
-                for b in heap[page_base..page_base + host_page_size()].iter_mut() {
-                    *b = 0x00;
-                }
             }
         }
 

lucet-runtime/lucet-runtime-internals/src/region/mmap.rs

@@ -9,6 +9,7 @@ license = "Apache-2.0 WITH LLVM-exception"
 cast = "0.2"
 clap = "2.23"
 failure = "0.1"
+human-size = "0.4"
 libc = "0.2"
 lucet-runtime = { path = "../lucet-runtime" }
 lucet-runtime-internals = { path = "../lucet-runtime/lucet-runtime-internals" }

lucet-wasi/Cargo.toml

@@ -2,7 +2,9 @@
 extern crate clap;
 
 use clap::Arg;
+use human_size::{Byte, Size};
 use lucet_runtime::{self, DlModule, Limits, MmapRegion, Module, Region};
+use lucet_runtime_internals::module::ModuleInternal;
 use lucet_wasi::{hostcalls, WasiCtxBuilder};
 use std::fs::File;
 use std::sync::Arc;
@@ -12,6 +14,7 @@ struct Config<'a> {
     guest_args: Vec<&'a str>,
     entrypoint: &'a str,
     preopen_dirs: Vec<(File, &'a str)>,
+    limits: Limits,
 }
 
 fn main() {
@@ -41,10 +44,10 @@ fn main() {
                      virtual filesystem. Each directory is specified as \
                      --dir `host_path:guest_path`, where `guest_path` specifies the path that will \
                      correspond to `host_path` for calls like `fopen` in the guest.\
-                     \
+                     \n\n\
                      For example, `--dir /home/host_user/wasi_sandbox:/sandbox` will make \
                      `/home/host_user/wasi_sandbox` available within the guest as `/sandbox`.\
-                     \
+                     \n\n\
                      Guests will be able to access any files and directories under the \
                      `host_path`, but will be unable to access other parts of the host \
                      filesystem through relative paths (e.g., `/sandbox/../some_other_file`) \
@@ -56,6 +59,27 @@ fn main() {
                 .required(true)
                 .help("Path to the `lucetc`-compiled WASI module"),
         )
+        .arg(
+            Arg::with_name("heap_memory_size")
+                .long("max-heap-size")
+                .takes_value(true)
+                .default_value("4 GiB")
+                .help("Maximum heap size (must be a multiple of 4 KiB)"),
+        )
+        .arg(
+            Arg::with_name("heap_address_space_size")
+                .long("heap-address-space")
+                .takes_value(true)
+                .default_value("8 GiB")
+                .help("Maximum heap address space size (must be a multiple of 4 KiB, and >= `max-heap-size`)"),
+        )
+        .arg(
+            Arg::with_name("stack_size")
+                .long("stack-size")
+                .takes_value(true)
+                .default_value("8 MiB")
+                .help("Maximum stack size (must be a multiple of 4 KiB)"),
+        )
         .arg(
             Arg::with_name("guest_args")
                 .required(false)
@@ -65,7 +89,9 @@ fn main() {
         .get_matches();
 
     let entrypoint = matches.value_of("entrypoint").unwrap();
+
     let lucet_module = matches.value_of("lucet_module").unwrap();
+
     let preopen_dirs = matches
         .values_of("preopen_dirs")
         .map(|vals| {
@@ -84,25 +110,66 @@ fn main() {
             .collect()
         })
         .unwrap_or(vec![]);
+
+    let heap_memory_size = value_t!(matches, "heap_memory_size", Size)
+        .unwrap_or_else(|e| e.exit())
+        .into::<Byte>()
+        .value() as usize;
+    let heap_address_space_size = value_t!(matches, "heap_address_space_size", Size)
+        .unwrap_or_else(|e| e.exit())
+        .into::<Byte>()
+        .value() as usize;
+
+    if heap_memory_size > heap_address_space_size {
+        println!("`heap-address-space` must be at least as large as `max-heap-size`");
+        println!("{}", matches.usage());
+        std::process::exit(1);
+    }
+
+    let stack_size = value_t!(matches, "stack_size", Size)
+        .unwrap_or_else(|e| e.exit())
+        .into::<Byte>()
+        .value() as usize;
+
+    let limits = Limits {
+        heap_memory_size,
+        heap_address_space_size,
+        stack_size,
+        globals_size: 0, // calculated from module
+    };
+
     let guest_args = matches
         .values_of("guest_args")
         .map(|vals| vals.collect())
         .unwrap_or(vec![]);
+
     let config = Config {
         lucet_module,
         guest_args,
         entrypoint,
         preopen_dirs,
+        limits,
     };
+
     run(config)
 }
 
 fn run(config: Config) {
     lucet_wasi::hostcalls::ensure_linked();
     let exitcode = {
         // doing all of this in a block makes sure everything gets dropped before exiting
-        let region = MmapRegion::create(1, &Limits::default()).expect("region can be created");
         let module = DlModule::load(&config.lucet_module).expect("module can be loaded");
+        let min_globals_size = module.globals().len() * std::mem::size_of::<u64>();
+        let globals_size = ((min_globals_size + 4096 - 1) / 4096) * 4096;
+
+        let region = MmapRegion::create(
+            1,
+            &Limits {
+                globals_size,
+                ..config.limits
+            },
+        )
+        .expect("region can be created");
 
         // put the path to the module on the front for argv[0]
         let args = std::iter::once(config.lucet_module)

lucet-wasi/src/main.rs

@@ -77,7 +77,7 @@ pub fn compile_sparse_page_data(compiler: &mut Compiler) -> Result<(), Error> {
     use crate::program::data::sparse::OwnedSparseData;
     let owned_data = OwnedSparseData::new(
         &compiler.prog.data_initializers()?,
-        compiler.prog.heap_spec(),
+        compiler.prog.heap_spec()?,
     );
     let sparse_data = owned_data.sparse_data();
 

lucetc/src/compiler/data.rs

@@ -76,7 +76,7 @@ impl<'p> EntityCreator<'p> {
         compiler: &Compiler,
     ) -> Result<ir::Heap, Error> {
         let base = self.bases.heap(func, compiler);
-        let heap_spec = self.program.heap_spec();
+        let heap_spec = self.program.heap_spec()?;
 
         self.cache.heap(index, || {
             if index != 0 {

lucetc/src/compiler/entity/mod.rs

@@ -5,7 +5,7 @@ use cranelift_module::{DataContext, Linkage};
 use failure::Error;
 
 pub fn compile_memory_specs(compiler: &mut Compiler) -> Result<(), Error> {
-    let heap = compiler.prog.heap_spec();
+    let heap = compiler.prog.heap_spec()?;
 
     let mut heap_spec_ctx = DataContext::new();
     heap_spec_ctx.define(serialize_spec(&heap).into_boxed_slice());

lucetc/src/compiler/memory.rs

@@ -7,10 +7,10 @@ use lucet_module_data::ModuleData;
 
 pub fn compile_module_data(compiler: &mut Compiler) -> Result<(), Error> {
     let module_data_serialized: Vec<u8> = {
-        let heap_spec = compiler.prog.heap_spec();
+        let heap_spec = compiler.prog.heap_spec()?;
         let compiled_data = OwnedSparseData::new(
             &compiler.prog.data_initializers()?,
-            compiler.prog.heap_spec(),
+            compiler.prog.heap_spec()?,
         );
         let sparse_data = compiled_data.sparse_data();
 

lucetc/src/compiler/module_data.rs

@@ -87,12 +87,20 @@ impl Lucetc {
         Ok(())
     }
 
-    pub fn reserved_size(mut self, reserved_size: u64) -> Self {
-        self.with_reserved_size(reserved_size);
+    pub fn min_reserved_size(mut self, min_reserved_size: u64) -> Self {
+        self.with_min_reserved_size(min_reserved_size);
         self
     }
-    pub fn with_reserved_size(&mut self, reserved_size: u64) {
-        self.heap.reserved_size = reserved_size;
+    pub fn with_min_reserved_size(&mut self, min_reserved_size: u64) {
+        self.heap.min_reserved_size = min_reserved_size;
+    }
+
+    pub fn max_reserved_size(mut self, max_reserved_size: u64) -> Self {
+        self.with_max_reserved_size(max_reserved_size);
+        self
+    }
+    pub fn with_max_reserved_size(&mut self, max_reserved_size: u64) {
+        self.heap.max_reserved_size = max_reserved_size;
     }
 
     pub fn guard_size(mut self, guard_size: u64) -> Self {

lucetc/src/lib.rs

@@ -50,8 +50,12 @@ pub fn run(opts: &Options) -> Result<(), Error> {
         c.with_builtins(builtins)?;
     }
 
-    if let Some(reserved_size) = opts.reserved_size {
-        c.with_reserved_size(reserved_size);
+    if let Some(min_reserved_size) = opts.min_reserved_size {
+        c.with_min_reserved_size(min_reserved_size);
+    }
+
+    if let Some(max_reserved_size) = opts.max_reserved_size {
+        c.with_max_reserved_size(max_reserved_size);
     }
 
     if let Some(guard_size) = opts.guard_size {