Replace stdout with a finalize function

api/examples/cart-checkout-validation-wasm-api.rs

@@ -38,8 +38,6 @@ fn main() -> Result<(), Box<dyn Error>> {
         1,
     )?;
 
-    context.finalize_output()?;
-
     Ok(())
 }
 

api/examples/echo.rs

@@ -11,8 +11,6 @@ fn main() -> Result<(), Box<dyn Error>> {
     let value = Value::deserialize(&input)?;
     let result = echo(value);
     result.serialize(&mut context)?;
-    context.finalize_output()?;
-
     Ok(())
 }
 

api/src/lib.rs

@@ -13,7 +13,6 @@
 //!     let input = context.input_get()?;
 //!     let value: i32 = Deserialize::deserialize(&input)?;
 //!     value.serialize(&mut context)?;
-//!     context.finalize_output()?;
 //!     Ok(())
 //! }
 //! ```
@@ -72,7 +71,6 @@ extern "C" {
     // Write API.
     fn shopify_function_output_new_bool(context: ContextPtr, bool: u32) -> usize;
     fn shopify_function_output_new_null(context: ContextPtr) -> usize;
-    fn shopify_function_output_finalize(context: ContextPtr) -> usize;
     fn shopify_function_output_new_i32(context: ContextPtr, int: i32) -> usize;
     fn shopify_function_output_new_f64(context: ContextPtr, float: f64) -> usize;
     fn shopify_function_output_new_utf8_str(
@@ -163,9 +161,6 @@ mod provider_fallback {
     pub(crate) unsafe fn shopify_function_output_new_null(context: ContextPtr) -> usize {
         shopify_function_provider::write::shopify_function_output_new_null(context) as usize
     }
-    pub(crate) unsafe fn shopify_function_output_finalize(context: ContextPtr) -> usize {
-        shopify_function_provider::write::shopify_function_output_finalize(context) as usize
-    }
     pub(crate) unsafe fn shopify_function_output_new_i32(context: ContextPtr, int: i32) -> usize {
         shopify_function_provider::write::shopify_function_output_new_i32(context, int) as usize
     }

api/src/shopify_function.h

@@ -122,15 +122,6 @@ __attribute__((import_module(SHOPIFY_FUNCTION_IMPORT_MODULE)))
 __attribute__((import_name("shopify_function_output_new_null")))
 extern WriteResult shopify_function_output_new_null(ContextPtr context);
 
-/**
- * Finalizes the output and returns the result
- * @param context The context pointer
- * @return WriteResult indicating success or failure
- */
-__attribute__((import_module(SHOPIFY_FUNCTION_IMPORT_MODULE)))
-__attribute__((import_name("shopify_function_output_finalize")))
-extern WriteResult shopify_function_output_finalize(ContextPtr context);
-
 /**
  * Creates a new 32-bit integer output value
  * @param context The context pointer

api/src/shopify_function.wat

@@ -151,18 +151,6 @@
     (func (param $context i32) (result i32))
   )
 
-  ;; Finalizes the output, making it available to the host.
-  ;; Must be called after output construction is complete.
-  ;; This is typically the last API call made before function returns.
-  ;; Signals that the response is complete and ready to be used.
-  ;; Parameters:
-  ;;   - context: i32 pointer to the context.
-  ;; Returns:
-  ;;   - i32 status code indicating success or failure
-  (import "shopify_function_v1" "shopify_function_output_finalize" 
-    (func (param $context i32) (result i32))
-  )
-
   ;; Writes a new integer output value.
   ;; Used for numeric values that fit within 32 bits.
   ;; More efficient than f64 for integral values.

api/src/test_data/header_test.c

@@ -18,7 +18,6 @@ volatile void* imports[] = {
     (void*)shopify_function_input_get_obj_key_at_index,
     (void*)shopify_function_output_new_bool,
     (void*)shopify_function_output_new_null,
-    (void*)shopify_function_output_finalize,
     (void*)shopify_function_output_new_i32,
     (void*)shopify_function_output_new_f64,
     (void*)shopify_function_output_new_utf8_str,
@@ -27,5 +26,4 @@ volatile void* imports[] = {
     (void*)shopify_function_output_finish_object,
     (void*)shopify_function_output_new_array,
     (void*)shopify_function_output_finish_array,
-    (void*)shopify_function_intern_utf8_str
-};
+    (void*)shopify_function_intern_utf8_str};

api/src/write.rs

@@ -115,11 +115,6 @@ impl Context {
         map_result(unsafe { crate::shopify_function_output_finish_array(self.0 as _) })
     }
 
-    /// Finalize the output. This must be called exactly once, and must be called after all other writes.
-    pub fn finalize_output(self) -> Result<(), Error> {
-        map_result(unsafe { crate::shopify_function_output_finalize(self.0 as _) })
-    }
-
     #[cfg(not(target_family = "wasm"))]
     /// Finalize the output and return the serialized value as a `serde_json::Value`.
     /// This is only available in non-Wasm targets, and therefore only recommended for use in tests.

integration_tests/tests/integration_test.rs

@@ -43,7 +43,7 @@ fn assert_fuel_consumed_within_threshold(target_fuel: u64, fuel_consumed: u64) {
     }
 }
 
-fn run_example(example: &str, input_bytes: Vec<u8>) -> Result<(Vec<u8>, u64)> {
+fn run_example(example: &str, input_bytes: Vec<u8>, is_wasi: bool) -> Result<(Vec<u8>, u64)> {
     let manifest_dir = env!("CARGO_MANIFEST_DIR");
     let workspace_root = std::path::PathBuf::from(manifest_dir).join("..");
     let engine = Engine::new(Config::new().consume_fuel(true))?;
@@ -97,6 +97,21 @@ fn run_example(example: &str, input_bytes: Vec<u8>) -> Result<(Vec<u8>, u64)> {
 
     let instructions = STARTING_FUEL.saturating_sub(store.get_fuel().unwrap_or_default());
 
+    let mut output = if !is_wasi && result.is_ok() {
+        let finalize_func = provider_instance.get_typed_func::<(), u64>(&mut store, "finalize")?;
+        let output = finalize_func.call(&mut store, ())?;
+        let ptr = (output >> u32::BITS) as u32;
+        let len = output as u32;
+        let mut output = vec![0; len as usize];
+        provider_instance
+            .get_memory(&mut store, "memory")
+            .unwrap()
+            .read(&store, ptr as usize, &mut output)?;
+        output
+    } else {
+        Vec::new()
+    };
+
     drop(store);
 
     if let Err(e) = result {
@@ -108,7 +123,9 @@ fn run_example(example: &str, input_bytes: Vec<u8>) -> Result<(Vec<u8>, u64)> {
         ));
     }
 
-    let output = stdout.contents().to_vec();
+    if is_wasi {
+        output = stdout.contents().to_vec();
+    }
     Ok((output, instructions))
 }
 
@@ -151,7 +168,7 @@ fn prepare_wasi_json_input(input: serde_json::Value) -> Result<Vec<u8>> {
 
 fn run_wasm_api_example(example: &str, input: serde_json::Value) -> Result<serde_json::Value> {
     let input_bytes = prepare_wasm_api_input(input)?;
-    let (output, _fuel) = run_example(example, input_bytes)?;
+    let (output, _fuel) = run_example(example, input_bytes, false)?;
     decode_msgpack_output(output)
 }
 
@@ -341,10 +358,11 @@ fn test_fuel_consumption_within_threshold() -> Result<()> {
         .map_err(|e| anyhow::anyhow!("Failed to prepare example: {}", e))?;
     let input = generate_cart_with_size(2, true);
     let wasm_api_input = prepare_wasm_api_input(input.clone())?;
-    let (_, wasm_api_fuel) = run_example("cart-checkout-validation-wasm-api", wasm_api_input)?;
+    let (_, wasm_api_fuel) =
+        run_example("cart-checkout-validation-wasm-api", wasm_api_input, false)?;
     eprintln!("WASM API fuel: {}", wasm_api_fuel);
     // Using a target fuel value as reference similar to the Javy example
-    assert_fuel_consumed_within_threshold(15880, wasm_api_fuel);
+    assert_fuel_consumed_within_threshold(13969, wasm_api_fuel);
     Ok(())
 }
 
@@ -361,12 +379,12 @@ fn test_benchmark_comparison_with_input() -> Result<()> {
 
     let wasm_api_input = prepare_wasm_api_input(input.clone())?;
     let (wasm_api_output, wasm_api_fuel) =
-        run_example("cart-checkout-validation-wasm-api", wasm_api_input)?;
+        run_example("cart-checkout-validation-wasm-api", wasm_api_input, false)?;
     let wasm_api_value = decode_msgpack_output(wasm_api_output)?;
 
     let wasi_json_input = prepare_wasi_json_input(input)?;
     let (non_wasm_api_output, non_wasm_api_fuel) =
-        run_example("cart-checkout-validation-wasi-json", wasi_json_input)?;
+        run_example("cart-checkout-validation-wasi-json", wasi_json_input, true)?;
     let non_wasm_api_value = decode_json_output(non_wasm_api_output)?;
 
     assert_eq!(wasm_api_value, non_wasm_api_value);
@@ -384,7 +402,7 @@ fn test_benchmark_comparison_with_input() -> Result<()> {
         wasm_api_fuel, non_wasm_api_fuel, improvement
     );
 
-    assert_fuel_consumed_within_threshold(15880, wasm_api_fuel);
+    assert_fuel_consumed_within_threshold(13969, wasm_api_fuel);
     assert_fuel_consumed_within_threshold(23858, non_wasm_api_fuel);
 
     Ok(())
@@ -403,12 +421,12 @@ fn test_benchmark_comparison_with_input_early_exit() -> Result<()> {
 
     let wasm_api_input = prepare_wasm_api_input(input.clone())?;
     let (wasm_api_output, wasm_api_fuel) =
-        run_example("cart-checkout-validation-wasm-api", wasm_api_input)?;
+        run_example("cart-checkout-validation-wasm-api", wasm_api_input, false)?;
     let wasm_api_value = decode_msgpack_output(wasm_api_output)?;
 
     let wasi_json_input = prepare_wasi_json_input(input)?;
     let (non_wasm_api_output, non_wasm_api_fuel) =
-        run_example("cart-checkout-validation-wasi-json", wasi_json_input)?;
+        run_example("cart-checkout-validation-wasi-json", wasi_json_input, true)?;
     let non_wasm_api_value = decode_json_output(non_wasm_api_output)?;
 
     assert_eq!(wasm_api_value, non_wasm_api_value);
@@ -427,7 +445,7 @@ fn test_benchmark_comparison_with_input_early_exit() -> Result<()> {
     );
 
     // Add fuel consumption threshold checks for both implementations
-    assert_fuel_consumed_within_threshold(17826, wasm_api_fuel);
+    assert_fuel_consumed_within_threshold(15534, wasm_api_fuel);
     assert_fuel_consumed_within_threshold(736695, non_wasm_api_fuel);
 
     Ok(())

provider/src/lib.rs

@@ -5,13 +5,22 @@ pub mod write;
 
 use rmp::encode::ByteBuf;
 use shopify_function_wasm_api_core::ContextPtr;
-use std::ptr::NonNull;
+use std::{cell::RefCell, ptr::NonNull};
 use string_interner::StringInterner;
 use write::State;
 
 pub const PROVIDER_MODULE_NAME: &str =
     concat!("shopify_function_v", env!("CARGO_PKG_VERSION_MAJOR"));
 
+// This serves 2 purposes:
+// 1. Creating multiple contexts in Wasm won't cause a hang waiting for input
+//    on stdin.
+// 2. Data in the context can be read by the host environment.
+// Thread local values live for the same amount of time as the Wasm instance.
+thread_local! {
+    static CONTEXT: RefCell<Option<ContextPtr>> = const { RefCell::new(None) };
+}
+
 #[cfg(target_pointer_width = "64")]
 type DoubleUsize = u128;
 #[cfg(target_pointer_width = "32")]
@@ -93,12 +102,22 @@ pub(crate) use decorate_for_target;
 #[cfg(target_family = "wasm")]
 #[export_name = "_shopify_function_context_new"]
 extern "C" fn shopify_function_context_new() -> ContextPtr {
-    Box::into_raw(Box::new(Context::new_from_stdin())) as _
+    // Assuming this is called from a Wasm module, we should return a
+    // pre-existing context if there is one to avoid trying to read input
+    // multiple times.
+    CONTEXT.with_borrow_mut(|option_ptr| {
+        *option_ptr.get_or_insert_with(|| Box::into_raw(Box::new(Context::new_from_stdin())) as _)
+    })
 }
 
 #[cfg(not(target_family = "wasm"))]
 pub fn shopify_function_context_new_from_msgpack_bytes(bytes: Vec<u8>) -> ContextPtr {
-    Box::into_raw(Box::new(Context::new(bytes))) as _
+    // Assuming this is called from a unit test, we should reset the context
+    // with the input passed as an argument instead of using any context that
+    // may have been present before.
+    CONTEXT.with_borrow_mut(|option_ptr| {
+        *option_ptr.insert(Box::into_raw(Box::new(Context::new(bytes))) as _)
+    })
 }
 
 decorate_for_target! {

provider/src/write.rs

@@ -1,7 +1,6 @@
 use crate::{decorate_for_target, Context, DoubleUsize};
 use rmp::encode;
 use shopify_function_wasm_api_core::{write::WriteResult, ContextPtr};
-use std::io::Write;
 
 mod state;
 
@@ -220,33 +219,6 @@ decorate_for_target! {
     }
 }
 
-decorate_for_target! {
-    fn shopify_function_output_finalize(context: ContextPtr) -> WriteResult {
-        match Context::mut_from_raw(context) {
-            Ok(context) => {
-                let Context {
-                    output_bytes,
-                    write_state,
-                    ..
-                } = &context;
-                if *write_state != State::End {
-                    return WriteResult::ValueNotFinished;
-                }
-                let mut stdout = std::io::stdout();
-                if stdout.write_all(output_bytes.as_slice()).is_err() {
-                    return WriteResult::IoError;
-                }
-                if stdout.flush().is_err() {
-                    return WriteResult::IoError;
-                }
-                let _ = unsafe { Box::from_raw(context as *mut Context) }; // drop the context
-                WriteResult::Ok
-            }
-            Err(_) => WriteResult::IoError,
-        }
-    }
-}
-
 #[cfg(not(target_family = "wasm"))]
 pub fn shopify_function_output_finalize_and_return_msgpack_bytes(
     context: ContextPtr,
@@ -269,6 +241,20 @@ pub fn shopify_function_output_finalize_and_return_msgpack_bytes(
     }
 }
 
+#[cfg(target_family = "wasm")]
+#[no_mangle]
+pub fn finalize() -> DoubleUsize {
+    use crate::CONTEXT;
+
+    CONTEXT.with_borrow(|context| {
+        let context = context.unwrap();
+        let context = Context::ref_from_raw(context).unwrap();
+        let output_ptr = context.output_bytes.as_vec().as_ptr();
+        let output_len = context.output_bytes.as_vec().len();
+        ((output_ptr as DoubleUsize) << usize::BITS) | output_len as DoubleUsize
+    })
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

trampoline/src/lib.rs

@@ -40,10 +40,6 @@ static IMPORTS: &[(&str, &str)] = &[
         "shopify_function_output_new_null",
         "_shopify_function_output_new_null",
     ),
-    (
-        "shopify_function_output_finalize",
-        "_shopify_function_output_finalize",
-    ),
     (
         "shopify_function_output_new_i32",
         "_shopify_function_output_new_i32",