Skip to content

Support WASM Api #460

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mssalemi merged 1 commit into from
May 7, 2025
Merged

Support WASM Api #460

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions build.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,4 +4,5 @@ fn main() {
println!("cargo:rerun-if-changed=providers/javy_quickjs_provider_v3.wasm");
println!("cargo:rerun-if-changed=providers/shopify_functions_javy_v1.wasm");
println!("cargo:rerun-if-changed=providers/shopify_functions_javy_v2.wasm");
println!("cargo:rerun-if-changed=providers/shopify_function_v1.wasm");
}
Binary file added providers/shopify_function_v1.wasm
View file
Open in desktop
Binary file not shown.
158 changes: 134 additions & 24 deletions src/engine.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,12 @@ pub struct ProfileOpts {
#[folder = "providers/"]
struct StandardProviders;

pub fn uses_msgpack_provider(module: &Module) -> bool {
module.imports().map(|i| i.module()).any(|module| {
module.starts_with("shopify_function_v") || module == "shopify_functions_javy_v2"
})
}

fn import_modules<T>(
module: &Module,
engine: &Engine,
Expand All @@ -28,8 +34,10 @@ fn import_modules<T>(
) {
let imported_modules: HashSet<String> =
module.imports().map(|i| i.module().to_string()).collect();

imported_modules.iter().for_each(|module_name| {
let imported_module_bytes = StandardProviders::get(&format!("{module_name}.wasm"));
let provider_path = format!("{module_name}.wasm");
let imported_module_bytes = StandardProviders::get(&provider_path);

if let Some(bytes) = imported_module_bytes {
let imported_module = Module::from_binary(engine, &bytes.data)
Expand All @@ -38,20 +46,22 @@ fn import_modules<T>(
let imported_module_instance = linker
.instantiate(&mut store, &imported_module)
.expect("Failed to instantiate imported instance");

linker
.instance(&mut store, module_name, imported_module_instance)
.expect("Failed to import module");
}
});
}

#[derive(Default)]
pub struct FunctionRunParams<'a> {
pub function_path: PathBuf,
pub input: BytesContainer,
pub export: &'a str,
pub profile_opts: Option<&'a ProfileOpts>,
pub scale_factor: f64,
pub module: Module,
pub engine: Engine,
}

const STARTING_FUEL: u64 = u64::MAX;
Expand Down Expand Up @@ -114,18 +124,10 @@ pub fn run(params: FunctionRunParams) -> Result<FunctionRunResult> {
export,
profile_opts,
scale_factor,
engine,
module,
} = params;

let engine = Engine::new(
Config::new()
.wasm_multi_memory(true)
.wasm_threads(false)
.consume_fuel(true)
.epoch_interruption(true),
)?;
let module = Module::from_file(&engine, &function_path)
.map_err(|e| anyhow!("Couldn't load the Function {:?}: {}", &function_path, e))?;

let input_stream = MemoryInputPipe::new(input.raw.clone());
let output_stream = MemoryOutputPipe::new(usize::MAX);
let error_stream = MemoryOutputPipe::new(usize::MAX);
Expand Down Expand Up @@ -201,11 +203,15 @@ pub fn run(params: FunctionRunParams) -> Result<FunctionRunResult> {

logs.extend_from_slice(error_logs.as_bytes());

let output_codec = input.codec;
let raw_output = output_stream
.try_into_inner()
.expect("Output stream reference still exists");

let output = BytesContainer::new(BytesContainerType::Output, input.codec, raw_output.to_vec())?;
let output = BytesContainer::new(
BytesContainerType::Output,
output_codec,
raw_output.to_vec(),
)?;

let name = function_path.file_name().unwrap().to_str().unwrap();
let size = function_path.metadata()?.len() / 1024;
Expand All @@ -226,6 +232,21 @@ pub fn run(params: FunctionRunParams) -> Result<FunctionRunResult> {
Ok(function_run_result)
}

/// Creates a new Engine with our standard configuration.
/// We use a dedicated function instead of making this the default configuration because:
/// 1. It's more explicit about what configuration we're using
/// 2. It keeps the door open for different configurations in the future without breaking changes
/// 3. It makes it easier to find all places where we create an Engine
pub fn new_engine() -> Result<Engine> {
Engine::new(
Config::new()
.wasm_multi_memory(true)
.wasm_threads(false)
.consume_fuel(true)
.epoch_interruption(true),
)
}

#[cfg(test)]
mod tests {
use colored::Colorize;
Expand All @@ -244,6 +265,9 @@ mod tests {

#[test]
fn test_js_function() -> Result<()> {
let engine = new_engine()?;
let module =
Module::from_file(&engine, Path::new("tests/fixtures/build/js_function.wasm"))?;
let input = json_input(include_bytes!(
"../tests/fixtures/input/js_function_input.json"
))?;
Expand All @@ -252,7 +276,10 @@ mod tests {
function_path: Path::new("tests/fixtures/build/js_function.wasm").to_path_buf(),
input,
export: DEFAULT_EXPORT,
..Default::default()
module,
engine,
scale_factor: 1.0,
profile_opts: None,
})?;

assert_eq!(function_run_result.memory_usage, 1280);
Expand All @@ -262,14 +289,22 @@ mod tests {

#[test]
fn test_js_v2_function() -> Result<()> {
let engine = new_engine()?;
let module = Module::from_file(
&engine,
Path::new("tests/fixtures/build/js_function_v2.wasm"),
)?;
let input = json_input(include_bytes!(
"../tests/fixtures/input/js_function_input.json"
))?;
let function_run_result = run(FunctionRunParams {
function_path: Path::new("tests/fixtures/build/js_function_v2.wasm").to_path_buf(),
input,
export: DEFAULT_EXPORT,
..Default::default()
module,
engine,
scale_factor: 1.0,
profile_opts: None,
})?;

assert_eq!(function_run_result.memory_usage, 1344);
Expand All @@ -278,6 +313,11 @@ mod tests {

#[test]
fn test_js_v3_function() -> Result<()> {
let engine = new_engine()?;
let module = Module::from_file(
&engine,
Path::new("tests/fixtures/build/js_function_v3.wasm"),
)?;
let input = json_input(include_bytes!(
"../tests/fixtures/input/js_function_input.json"
))?;
Expand All @@ -286,7 +326,10 @@ mod tests {
function_path: Path::new("tests/fixtures/build/js_function_v3.wasm").to_path_buf(),
input,
export: DEFAULT_EXPORT,
..Default::default()
module,
engine,
scale_factor: 1.0,
profile_opts: None,
})?;

assert_eq!(function_run_result.memory_usage, 1344);
Expand All @@ -295,6 +338,11 @@ mod tests {

#[test]
fn test_js_functions_javy_v1() -> Result<()> {
let engine = new_engine()?;
let module = Module::from_file(
&engine,
Path::new("tests/fixtures/build/js_functions_javy_v1.wasm"),
)?;
let input = json_input(include_bytes!(
"../tests/fixtures/input/js_function_input.json"
))?;
Expand All @@ -304,7 +352,10 @@ mod tests {
.to_path_buf(),
input,
export: DEFAULT_EXPORT,
..Default::default()
module,
engine,
scale_factor: 1.0,
profile_opts: None,
})?;

assert_eq!(function_run_result.memory_usage, 1344);
Expand All @@ -313,11 +364,16 @@ mod tests {

#[test]
fn test_exit_code_zero() -> Result<()> {
let engine = new_engine()?;
let module = Module::from_file(&engine, Path::new("tests/fixtures/build/exit_code.wasm"))?;
let function_run_result = run(FunctionRunParams {
function_path: Path::new("tests/fixtures/build/exit_code.wasm").to_path_buf(),
input: json_input(&serde_json::to_vec(&json!({ "code": 0 }))?)?,
export: DEFAULT_EXPORT,
..Default::default()
module,
engine,
scale_factor: 1.0,
profile_opts: None,
})?;

assert_eq!(function_run_result.logs, "");
Expand All @@ -326,11 +382,16 @@ mod tests {

#[test]
fn test_exit_code_one() -> Result<()> {
let engine = new_engine()?;
let module = Module::from_file(&engine, Path::new("tests/fixtures/build/exit_code.wasm"))?;
let function_run_result = run(FunctionRunParams {
function_path: Path::new("tests/fixtures/build/exit_code.wasm").to_path_buf(),
input: json_input(&serde_json::to_vec(&json!({ "code": 1 }))?)?,
export: DEFAULT_EXPORT,
..Default::default()
module,
engine,
scale_factor: 1.0,
profile_opts: None,
})?;

assert_eq!(function_run_result.logs, "module exited with code: 1");
Expand All @@ -339,11 +400,19 @@ mod tests {

#[test]
fn test_linear_memory_usage_in_kb() -> Result<()> {
let engine = new_engine()?;
let module = Module::from_file(
&engine,
Path::new("tests/fixtures/build/linear_memory.wasm"),
)?;
let function_run_result = run(FunctionRunParams {
function_path: Path::new("tests/fixtures/build/linear_memory.wasm").to_path_buf(),
input: json_input(&serde_json::to_vec(&json!({}))?)?,
export: DEFAULT_EXPORT,
..Default::default()
module,
engine,
scale_factor: 1.0,
profile_opts: None,
})?;

assert_eq!(function_run_result.memory_usage, 12800); // 200 * 64KiB pages
Expand All @@ -352,12 +421,20 @@ mod tests {

#[test]
fn test_logs_truncation() -> Result<()> {
let engine = new_engine()?;
let module = Module::from_file(
&engine,
Path::new("tests/fixtures/build/log_truncation_function.wasm"),
)?;
let function_run_result = run(FunctionRunParams {
input: json_input("{}".as_bytes())?,
function_path: Path::new("tests/fixtures/build/log_truncation_function.wasm")
.to_path_buf(),
export: DEFAULT_EXPORT,
..Default::default()
module,
engine,
scale_factor: 1.0,
profile_opts: None,
})?;

assert!(
Expand All @@ -374,12 +451,16 @@ mod tests {
#[test]
fn test_file_size_in_kb() -> Result<()> {
let file_path = Path::new("tests/fixtures/build/exit_code.wasm");

let engine = new_engine()?;
let module = Module::from_file(&engine, file_path)?;
let function_run_result = run(FunctionRunParams {
function_path: file_path.to_path_buf(),
input: json_input(&serde_json::to_vec(&json!({ "code": 0 }))?)?,
export: DEFAULT_EXPORT,
..Default::default()
module,
engine,
scale_factor: 1.0,
profile_opts: None,
})?;

assert_eq!(
Expand All @@ -388,4 +469,33 @@ mod tests {
);
Ok(())
}

#[test]
fn test_wasm_api_function() -> Result<()> {
let engine = new_engine()?;
let module = Module::from_file(
&engine,
Path::new("tests/fixtures/build/echo.trampolined.wasm"),
)?;
let expected_input_value = json!({"foo": "echo", "bar": "test"});
let input = serde_json::to_vec(&expected_input_value).unwrap();
let input_bytes = BytesContainer::new(BytesContainerType::Input, Codec::Json, input);
let function_run_result = run(FunctionRunParams {
function_path: Path::new("tests/fixtures/build/echo.trampolined.wasm").to_path_buf(),
input: input_bytes.unwrap(),
export: DEFAULT_EXPORT,
module,
engine,
scale_factor: 1.0,
profile_opts: None,
});

assert!(function_run_result.is_ok());
let result = function_run_result.unwrap();
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&result.input.raw).unwrap(),
expected_input_value
);
Ok(())
}
}
19 changes: 15 additions & 4 deletions src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use function_runner::{BytesContainer, BytesContainerType, Codec};
use wasmtime::Module;

use std::{
fs::File,
Expand Down Expand Up @@ -53,9 +54,6 @@ struct Opts {
#[clap(long)]
profile_frequency: Option<u32>,

#[clap(short = 'c', long, value_enum, default_value = "json")]
codec: Codec,

/// Path to graphql file containing Function schema; if omitted, defaults will be used to calculate limits.
#[clap(short = 's', long)]
schema_path: Option<PathBuf>,
Expand Down Expand Up @@ -135,7 +133,18 @@ fn main() -> Result<()> {

let query_string = opts.read_query_to_string().transpose()?;

let input = BytesContainer::new(BytesContainerType::Input, opts.codec, buffer)?;
let engine = function_runner::engine::new_engine()?;
let module = Module::from_file(&engine, &opts.function)
.map_err(|e| anyhow!("Couldn't load the Function {:?}: {}", &opts.function, e))?;

// Infer codec from the module based on imported modules
let codec = if function_runner::engine::uses_msgpack_provider(&module) {
Codec::Messagepack
} else {
Codec::Json
};

let input = BytesContainer::new(BytesContainerType::Input, codec, buffer)?;
let scale_factor = if let (Some(schema_string), Some(query_string), Some(json_value)) =
(schema_string, query_string, input.json_value.clone())
{
Expand All @@ -158,6 +167,8 @@ fn main() -> Result<()> {
export: opts.export.as_ref(),
profile_opts: profile_opts.as_ref(),
scale_factor,
module,
engine,
})?;

if opts.json {
Expand Down
Binary file added tests/fixtures/build/echo.trampolined.wasm
View file
Open in desktop
Binary file not shown.
Loading