Implement memory usage logging

src/engine.rs

@@ -26,6 +26,7 @@ pub fn run(script_path: PathBuf, input_path: PathBuf) -> Result<FunctionRunResul
     let error_stream = wasi_common::pipe::WritePipe::new_in_memory();
 
     let runtime: Duration;
+    let memory_usage: u64;
 
     {
         let mut linker = Linker::new(&engine);
@@ -43,11 +44,17 @@ pub fn run(script_path: PathBuf, input_path: PathBuf) -> Result<FunctionRunResul
         let start = Instant::now();
 
         let instance = linker.instantiate(&mut store, &module)?;
+
+        let memory = instance
+            .get_memory(&mut store, "memory")
+            .ok_or(anyhow::format_err!("failed to find `memory` export"))?;
+
         let module_result = instance
             .get_typed_func::<(), (), _>(&mut store, "_start")?
             .call(&mut store, ());
 
         runtime = start.elapsed();
+        memory_usage = memory.size(&store);
 
         match module_result {
             Ok(_) => {}
@@ -71,7 +78,8 @@ pub fn run(script_path: PathBuf, input_path: PathBuf) -> Result<FunctionRunResul
     let output: serde_json::Value = serde_json::from_slice(output.as_slice())
         .map_err(|e| anyhow!("Couldn't decode Script Output: {}", e))?;
 
-    let function_run_result = FunctionRunResult::new(runtime, output, logs.to_string());
+    let function_run_result =
+        FunctionRunResult::new(runtime, memory_usage, output, logs.to_string());
 
     Ok(function_run_result)
 }
@@ -81,6 +89,9 @@ mod tests {
     use super::*;
     use std::path::Path;
 
+    // Arbitrary, used to verify that the runner works as expected.
+    const RUNTIME_THRESHOLD: Duration = Duration::from_millis(5);
+
     #[test]
     fn test_runtime_under_threshold() {
         let function_run_result = run(
@@ -89,7 +100,7 @@ mod tests {
         )
         .unwrap();
 
-        assert!(function_run_result.runtime <= Duration::from_millis(5));
+        assert!(function_run_result.runtime <= RUNTIME_THRESHOLD);
     }
 
     #[test]
@@ -100,6 +111,38 @@ mod tests {
         )
         .unwrap();
 
-        assert!(function_run_result.runtime > Duration::from_millis(5));
+        assert!(function_run_result.runtime > RUNTIME_THRESHOLD);
+    }
+
+    #[test]
+    fn test_memory_usage_under_threshold() {
+        let function_run_result = run(
+            Path::new("tests/benchmarks/hello_world.wasm").to_path_buf(),
+            Path::new("tests/benchmarks/hello_world.json").to_path_buf(),
+        )
+        .unwrap();
+
+        assert_eq!(function_run_result.memory_usage, 17);
+    }
+
+    #[test]
+    fn test_memory_usage_over_threshold() {
+        let function_run_result = run(
+            Path::new("tests/benchmarks/hello_42_pages.wasm").to_path_buf(),
+            Path::new("tests/benchmarks/hello_42_pages.json").to_path_buf(),
+        )
+        .unwrap();
+
+        assert_eq!(function_run_result.memory_usage, 42);
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_panic() {
+        run(
+            Path::new("tests/benchmarks/stack_overflow.wasm").to_path_buf(),
+            Path::new("tests/benchmarks/stack_overflow.json").to_path_buf(),
+        )
+        .unwrap();
     }
 }

src/function_run_result.rs

@@ -3,14 +3,21 @@ use std::{fmt, time::Duration};
 
 pub struct FunctionRunResult {
     pub runtime: Duration,
+    pub memory_usage: u64,
     pub logs: String,
     pub output: serde_json::Value,
 }
 
 impl FunctionRunResult {
-    pub fn new(runtime: Duration, output: serde_json::Value, logs: String) -> Self {
+    pub fn new(
+        runtime: Duration,
+        memory_usage: u64,
+        output: serde_json::Value,
+        logs: String,
+    ) -> Self {
         FunctionRunResult {
             runtime,
+            memory_usage,
             output,
             logs,
         }
@@ -22,7 +29,8 @@ impl fmt::Display for FunctionRunResult {
         let title = "      Benchmark Results      ".black().on_bright_green();
         write!(f, "{}\n\n", title)?;
 
-        writeln!(f, "Runtime: {:?}\n", self.runtime)?;
+        writeln!(f, "Runtime: {:?}", self.runtime)?;
+        writeln!(f, "Memory Usage: {}KB\n", self.memory_usage * 64)?;
 
         writeln!(
             f,

tests/benchmarks/hello_42_pages.json

@@ -0,0 +1,11 @@
+{
+  "input": {
+    "context": {
+      "suffix": "From Core!"
+    }
+  },
+  "configuration": {
+    "message": "From Core!"
+  },
+  "extension_point": "hello_world"
+}

tests/benchmarks/stack_overflow.json

@@ -0,0 +1,11 @@
+{
+  "input": {
+    "context": {
+      "suffix": "From Core!"
+    }
+  },
+  "configuration": {
+    "message": "From Core!"
+  },
+  "extension_point": "hello_world"
+}