Merge pull request #15 from Shopify/reformat-move-main-logic

refactor: move engine function outside of the main

Author: hainsdominic

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "script-runner"
-version = "0.2.1"
+version = "0.2.2"
 edition = "2021"
 
 [profile.test]

Cargo.toml

@@ -0,0 +1,105 @@
+use anyhow::{anyhow, Result};
+use std::{
+    path::PathBuf,
+    time::{Duration, Instant},
+};
+use wasmtime::{Engine, Linker, Module, Store};
+
+use wasmtime_wasi::sync::WasiCtxBuilder;
+
+use crate::function_run_result::FunctionRunResult;
+
+pub fn run(script_path: PathBuf, input_path: PathBuf) -> Result<FunctionRunResult> {
+    let engine = Engine::default();
+    let module = Module::from_file(&engine, &script_path)
+        .map_err(|e| anyhow!("Couldn't load script {:?}: {}", &script_path, e))?;
+
+    let input: serde_json::Value = serde_json::from_reader(
+        std::fs::File::open(&input_path)
+            .map_err(|e| anyhow!("Couldn't load input {:?}: {}", &input_path, e))?,
+    )
+    .map_err(|e| anyhow!("Couldn't load input {:?}: {}", &input_path, e))?;
+    let input = serde_json::to_vec(&input)?;
+
+    let input_stream = wasi_common::pipe::ReadPipe::new(std::io::Cursor::new(input));
+    let output_stream = wasi_common::pipe::WritePipe::new_in_memory();
+    let error_stream = wasi_common::pipe::WritePipe::new_in_memory();
+
+    let runtime: Duration;
+
+    {
+        let mut linker = Linker::new(&engine);
+        wasmtime_wasi::add_to_linker(&mut linker, |s| s)?;
+        let wasi = WasiCtxBuilder::new()
+            .stdin(Box::new(input_stream))
+            .stdout(Box::new(output_stream.clone()))
+            .stderr(Box::new(error_stream.clone()))
+            .inherit_args()?
+            .build();
+        let mut store = Store::new(&engine, wasi);
+
+        linker.module(&mut store, "", &module)?;
+
+        let start = Instant::now();
+
+        let instance = linker.instantiate(&mut store, &module)?;
+        let module_result = instance
+            .get_typed_func::<(), (), _>(&mut store, "_start")?
+            .call(&mut store, ());
+
+        runtime = start.elapsed();
+
+        match module_result {
+            Ok(_) => {}
+            Err(e) => {
+                eprintln!("Error:\n{}", e);
+            }
+        }
+    };
+
+    let logs = error_stream
+        .try_into_inner()
+        .expect("Error stream reference still exists")
+        .into_inner();
+    let logs =
+        std::str::from_utf8(&logs).map_err(|e| anyhow!("Couldn't print Script Logs: {}", e))?;
+
+    let output = output_stream
+        .try_into_inner()
+        .expect("Output stream reference still exists")
+        .into_inner();
+    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());
+
+    Ok(function_run_result)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::path::Path;
+
+    #[test]
+    fn test_runtime_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!(function_run_result.runtime <= Duration::from_millis(5));
+    }
+
+    #[test]
+    fn test_runtime_over_threshold() {
+        let function_run_result = run(
+            Path::new("tests/benchmarks/sleeps.wasm").to_path_buf(),
+            Path::new("tests/benchmarks/sleeps.json").to_path_buf(),
+        )
+        .unwrap();
+
+        assert!(function_run_result.runtime > Duration::from_millis(5));
+    }
+}

src/engine.rs

@@ -0,0 +1,42 @@
+use colored::Colorize;
+use std::{fmt, time::Duration};
+
+pub struct FunctionRunResult {
+    pub runtime: Duration,
+    pub logs: String,
+    pub output: serde_json::Value,
+}
+
+impl FunctionRunResult {
+    pub fn new(runtime: Duration, output: serde_json::Value, logs: String) -> Self {
+        FunctionRunResult {
+            runtime,
+            output,
+            logs,
+        }
+    }
+}
+
+impl fmt::Display for FunctionRunResult {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let title = "      Benchmark Results      ".black().on_bright_green();
+        write!(f, "{}\n\n", title)?;
+
+        writeln!(f, "Runtime: {:?}\n", self.runtime)?;
+
+        writeln!(
+            f,
+            "{}\n\n{}",
+            "            Logs             ".black().on_bright_blue(),
+            self.logs
+        )?;
+
+        writeln!(
+            f,
+            "Output:\n{}",
+            serde_json::to_string_pretty(&self.output).unwrap_or_else(|error| error.to_string())
+        )?;
+
+        Ok(())
+    }
+}

src/function_run_result.rs

@@ -1 +1,2 @@
-pub mod run_statistics;
+pub mod engine;
+pub mod function_run_result;

src/lib.rs

@@ -1,14 +1,8 @@
-use std::{
-    path::PathBuf,
-    time::{Duration, Instant},
-};
+use std::path::PathBuf;
 
-use anyhow::{anyhow, Result};
+use anyhow::Result;
 use clap::Parser;
-use colored::Colorize;
-use script_runner::run_statistics::RunStatistics;
-use wasmtime::*;
-use wasmtime_wasi::sync::WasiCtxBuilder;
+use script_runner::engine::run;
 
 /// Simple script runner which takes JSON as a convenience.
 #[derive(Parser)]
@@ -25,77 +19,9 @@ struct Opts {
 fn main() -> Result<()> {
     let opts: Opts = Opts::parse();
 
-    let engine = Engine::default();
-    let module = Module::from_file(&engine, &opts.script)
-        .map_err(|e| anyhow!("Couldn't load script {:?}: {}", &opts.script, e))?;
+    let function_run_result = run(opts.script, opts.input)?;
 
-    let input: serde_json::Value = serde_json::from_reader(
-        std::fs::File::open(&opts.input)
-            .map_err(|e| anyhow!("Couldn't load input {:?}: {}", &opts.input, e))?,
-    )
-    .map_err(|e| anyhow!("Couldn't load input {:?}: {}", &opts.input, e))?;
-    let input = serde_json::to_vec(&input)?;
-
-    let input_stream = wasi_common::pipe::ReadPipe::new(std::io::Cursor::new(input));
-    let output_stream = wasi_common::pipe::WritePipe::new_in_memory();
-    let error_stream = wasi_common::pipe::WritePipe::new_in_memory();
-
-    {
-        // Link WASI and construct the store.
-        let mut linker = Linker::new(&engine);
-        wasmtime_wasi::add_to_linker(&mut linker, |s| s)?;
-        let wasi = WasiCtxBuilder::new()
-            .stdin(Box::new(input_stream))
-            .stdout(Box::new(output_stream.clone()))
-            .stderr(Box::new(error_stream.clone()))
-            .inherit_args()?
-            .build();
-        let mut store = Store::new(&engine, wasi);
-
-        linker.module(&mut store, "", &module)?;
-
-        let start = Instant::now();
-
-        // Execute the module
-        let result = linker
-            .get_default(&mut store, "")?
-            .typed::<(), (), _>(&store)?
-            .call(&mut store, ());
-
-        let runtime = start.elapsed();
-
-        let benchmark = RunStatistics::new(runtime, Duration::from_millis(5));
-
-        println!("{}", benchmark);
-
-        match result {
-            Ok(_) => {}
-            Err(e) => {
-                eprintln!("Error:\n{}", e);
-            }
-        }
-    };
-
-    let logs = error_stream
-        .try_into_inner()
-        .expect("Error stream reference still exists")
-        .into_inner();
-    let logs =
-        std::str::from_utf8(&logs).map_err(|e| anyhow!("Couldn't print Script Logs: {}", e))?;
-
-    println!(
-        "{}\n\n{}",
-        "            Logs             ".black().on_bright_blue(),
-        logs
-    );
-
-    let output = output_stream
-        .try_into_inner()
-        .expect("Output stream reference still exists")
-        .into_inner();
-    let output: serde_json::Value = serde_json::from_slice(output.as_slice())
-        .map_err(|e| anyhow!("Couldn't decode Script Output: {}", e))?;
-    println!("Output:\n{}", serde_json::to_string_pretty(&output)?);
+    println!("{}", function_run_result);
 
     Ok(())
 }