From 3375c7a3b6b029ab5859ff00c1554abc8597542b Mon Sep 17 00:00:00 2001
From: Valery Antopol <valery.antopol@gmail.com>
Date: Wed, 24 Jan 2024 16:16:33 +0300
Subject: [PATCH] feat(deps)!: update to marine runtime with memory limits and
 wasmtime  (#768)

---
 air/Cargo.toml                                |  2 +-
 .../tetraplets/security_tetraplets.rs         |  1 +
 avm/server/Cargo.toml                         |  2 +-
 avm/server/src/runner.rs                      | 24 ++++++++++++-------
 4 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/air/Cargo.toml b/air/Cargo.toml
index 75446b4f..829d4198 100644
--- a/air/Cargo.toml
+++ b/air/Cargo.toml
@@ -52,7 +52,7 @@ rkyv = { version = "0.7.43", features = ["strict", "validation", "size_32", "arc
 [dev_dependencies]
 air-test-utils = { path = "../crates/air-lib/test-utils" }
 air-testing-framework = { path = "../crates/testing-framework" }
-fluence-app-service = "0.29.0"
+fluence-app-service = "0.31.0"
 marine-rs-sdk = { version = "0.10.0", features = ["logger"] }
 
 borsh = "0.10.3"
diff --git a/air/tests/test_module/features/tetraplets/security_tetraplets.rs b/air/tests/test_module/features/tetraplets/security_tetraplets.rs
index 47c2eb70..d8fd01fe 100644
--- a/air/tests/test_module/features/tetraplets/security_tetraplets.rs
+++ b/air/tests/test_module/features/tetraplets/security_tetraplets.rs
@@ -317,6 +317,7 @@ fn construct_service_config(module_name: impl Into<String>) -> AppServiceConfig
 
     let marine_config = MarineConfig {
         modules_dir: Some(PathBuf::from(module_path)),
+        total_memory_limit: None,
         modules_config: vec![module_descriptor],
         default_modules_config: None,
     };
diff --git a/avm/server/Cargo.toml b/avm/server/Cargo.toml
index 449d0c1d..edbcab3a 100644
--- a/avm/server/Cargo.toml
+++ b/avm/server/Cargo.toml
@@ -19,7 +19,7 @@ air-interpreter-interface = { version = "0.17.1", path = "../../crates/air-lib/i
 air-interpreter-sede = { version = "0.1.0", path = "../../crates/air-lib/interpreter-sede" }
 air-utils = { version = "0.2.0", path = "../../crates/air-lib/utils" }
 avm-data-store = { version = "0.7.6", path = "../../crates/data-store" }
-marine-runtime = "0.30.0"
+marine-runtime = "0.32.0"
 polyplets = { version = "0.6.0", path = "../../crates/air-lib/polyplets" }
 avm-interface = { version = "0.31.1", path = "../../avm/interface" }
 
diff --git a/avm/server/src/runner.rs b/avm/server/src/runner.rs
index 15b3f6cd..c60407ef 100644
--- a/avm/server/src/runner.rs
+++ b/avm/server/src/runner.rs
@@ -36,6 +36,8 @@ pub struct AVMRunner {
     marine: Marine,
     /// file name of the AIR interpreter .wasm
     wasm_filename: String,
+    /// The memory limit provided by constructor
+    total_memory_limit: Option<u64>,
 }
 
 /// Return statistic of AVM server Wasm module heap footprint.
@@ -44,25 +46,29 @@ pub struct AVMMemoryStats {
     /// Please note that linear memory contains not only heap, but globals, shadow stack and so on.
     pub memory_size: usize,
     /// Possibly set max memory size for AVM server.
-    pub max_memory_size: Option<usize>,
+    pub total_memory_limit: Option<u64>,
+    /// Number of allocations rejected due to memory limit.
+    /// May be not recorded by some backends in Marine.
+    pub allocation_rejects: Option<u32>,
 }
 
 impl AVMRunner {
     /// Create AVM with the provided config.
     pub fn new(
         air_wasm_path: PathBuf,
-        max_heap_size: Option<u64>,
+        total_memory_limit: Option<u64>,
         logging_mask: i32,
     ) -> RunnerResult<Self> {
         let (wasm_dir, wasm_filename) = split_dirname(air_wasm_path)?;
 
         let marine_config =
-            make_marine_config(wasm_dir, &wasm_filename, max_heap_size, logging_mask);
+            make_marine_config(wasm_dir, &wasm_filename, total_memory_limit, logging_mask);
         let marine = Marine::with_raw_config(marine_config)?;
 
         let avm = Self {
             marine,
             wasm_filename,
+            total_memory_limit,
         };
 
         Ok(avm)
@@ -190,11 +196,12 @@ impl AVMRunner {
         let stats = self.marine.module_memory_stats();
 
         // only the interpreters must be loaded in Marine
-        debug_assert!(stats.len() == 1);
+        debug_assert!(stats.modules.len() == 1);
 
         AVMMemoryStats {
-            memory_size: stats[0].memory_size,
-            max_memory_size: stats[0].max_memory_size,
+            memory_size: stats.modules[0].memory_size,
+            total_memory_limit: self.total_memory_limit,
+            allocation_rejects: stats.allocation_stats.map(|stats| stats.allocation_rejects),
         }
     }
 }
@@ -279,12 +286,10 @@ fn split_dirname(path: PathBuf) -> RunnerResult<(PathBuf, String)> {
 fn make_marine_config(
     air_wasm_dir: PathBuf,
     air_wasm_file: &str,
-    max_heap_size: Option<u64>,
+    total_memory_limit: Option<u64>,
     logging_mask: i32,
 ) -> MarineConfig {
     let air_module_config = marine::MarineModuleConfig {
-        mem_pages_count: None,
-        max_heap_size,
         logger_enabled: true,
         host_imports: <_>::default(),
         wasi: None,
@@ -293,6 +298,7 @@ fn make_marine_config(
 
     MarineConfig {
         modules_dir: Some(air_wasm_dir),
+        total_memory_limit,
         modules_config: vec![ModuleDescriptor {
             load_from: None,
             file_name: String::from(air_wasm_file),