feat!: add memory limiting

Cargo.toml

@@ -32,6 +32,7 @@ members = [
     "marine",
     "marine/tests/wasm_tests/arguments_passing",
     "marine/tests/wasm_tests/arrays_passing",
+    "marine/tests/wasm_tests/memory_limiting",
     "marine/tests/wasm_tests/records_passing",
     "marine/tests/wasm_tests/wasi",
     "marine-js",

core/src/config.rs

@@ -25,9 +25,6 @@ use std::path::PathBuf;
 use std::collections::HashMap;
 use std::collections::HashSet;
 
-// 65536*1600 ~ 100 Mb (Wasm page size is 64 Kb)
-const DEFAULT_HEAP_PAGES_COUNT: u32 = 1600;
-
 pub type ErrorHandler =
     Option<Box<dyn Fn(&HostImportError) -> Option<IValue> + Sync + Send + 'static>>;
 pub type HostExportedFunc<WB> = Box<
@@ -56,10 +53,6 @@ pub struct HostImportDescriptor<WB: WasmBackend> {
 }
 
 pub struct MModuleConfig<WB: WasmBackend> {
-    /// Maximum number of Wasm memory pages that loaded module can use.
-    /// Each Wasm page is 65536 bytes long.
-    pub max_heap_pages_count: u32,
-
     /// Import object that will be used in module instantiation process.
     pub raw_imports: HashMap<String, RawImportCreator<WB>>,
 
@@ -74,7 +67,6 @@ impl<WB: WasmBackend> Default for MModuleConfig<WB> {
     fn default() -> Self {
         // some reasonable defaults
         Self {
-            max_heap_pages_count: DEFAULT_HEAP_PAGES_COUNT,
             raw_imports: HashMap::new(),
             host_imports: HashMap::new(),
             wasi_parameters: WasiParameters::default(),
@@ -86,11 +78,6 @@ impl<WB: WasmBackend> Default for MModuleConfig<WB> {
 
 #[allow(dead_code)]
 impl<WB: WasmBackend> MModuleConfig<WB> {
-    pub fn with_mem_pages_count(mut self, mem_pages_count: u32) -> Self {
-        self.max_heap_pages_count = mem_pages_count;
-        self
-    }
-
     pub fn with_wasi_envs(mut self, envs: HashMap<String, String>) -> Self {
         self.wasi_parameters.envs = envs;
         self
@@ -106,3 +93,39 @@ impl<WB: WasmBackend> MModuleConfig<WB> {
         self
     }
 }
+
+pub struct MarineCoreConfig {
+    pub(crate) memory_limit: u64,
+}
+
+impl Default for MarineCoreConfig {
+    fn default() -> Self {
+        Self {
+            memory_limit: DEFAULT_MEMORY_LIMIT,
+        }
+    }
+}
+
+pub const DEFAULT_MEMORY_LIMIT: u64 = 0xFFFFFFFFFFFFFFFF;
+
+#[derive(Default, Debug)]
+pub struct MarineCoreConfigBuilder {
+    memory_limit: Option<u64>,
+}
+
+impl MarineCoreConfigBuilder {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn memory_limit(mut self, memory_limit: u64) -> Self {
+        self.memory_limit = Some(memory_limit);
+        self
+    }
+
+    pub fn build(self) -> MarineCoreConfig {
+        MarineCoreConfig {
+            memory_limit: self.memory_limit.unwrap_or(DEFAULT_MEMORY_LIMIT),
+        }
+    }
+}

core/src/lib.rs

@@ -36,6 +36,9 @@ mod module;
 mod memory_statistic;
 
 pub use crate::marine_core::MModuleInterface;
+pub use config::MarineCoreConfig;
+pub use config::MarineCoreConfigBuilder;
+pub use config::DEFAULT_MEMORY_LIMIT;
 pub use errors::MError;
 pub use host_imports::HostImportError;
 pub use module::IValue;

core/src/marine_core.rs

@@ -15,6 +15,7 @@
  */
 
 use super::generic::*;
+use crate::config::MarineCoreConfig;
 use crate::module::MModule;
 use crate::module::MRecordTypes;
 use crate::{IRecordType, IValue, MemoryStats, MError, MFunctionSignature, ModuleMemoryStat, MResult};
@@ -60,9 +61,10 @@ pub struct MarineCore<WB: WasmBackend> {
 }
 
 impl<WB: WasmBackend> MarineCore<WB> {
-    pub fn new() -> MResult<Self> {
+    pub fn new(config: MarineCoreConfig) -> MResult<Self> {
         let wasm_backend = WB::new()?;
-        let store = <WB as WasmBackend>::Store::new(&wasm_backend);
+        let mut store = <WB as WasmBackend>::Store::new(&wasm_backend);
+        store.set_memory_limit(config.memory_limit);
         Ok(Self {
             modules: HashMap::new(),
             wasm_backend,
@@ -184,12 +186,15 @@ impl<WB: WasmBackend> MarineCore<WB> {
                 ModuleMemoryStat::new(
                     module_name,
                     module.memory_size(&mut self.store.borrow_mut().as_context_mut()),
-                    module.max_memory_size(),
                 )
             })
             .collect::<Vec<_>>();
+        let allocation_stats = self.store.borrow_mut().report_memory_allocation_stats();
+        MemoryStats::new(records, allocation_stats)
+    }
 
-        records.into()
+    pub fn clear_allocation_stats(&mut self) {
+        self.store.borrow_mut().clear_allocation_stats()
     }
 
     fn get_module_interface(module: &MModule<WB>) -> MModuleInterface<'_> {