WASM/WASI threads support

dotnet/nuget/Extism.runtime.win.csproj

@@ -7,10 +7,11 @@
     </PropertyGroup>
 
     <PropertyGroup>
-        <PackageId>Extism.runtime.win-x64</PackageId>
-        <Version>0.6.0</Version>
+        <!-- TODO: Temporal package name. Real one: Extism.runtime.win-x64 -->
+        <PackageId>Extism.runtime.wasm-threads.win-x64</PackageId>
+        <Version>0.6.0-threads01</Version>
         <Authors>Extism Contributors</Authors>
-        <Description>Internal implementation package for Extism to work on Windows x64</Description>
+        <Description>Fork of the original Windows Extism 0.6.0 runtime with WASM threads feature enabled for Wasmtime, when EXTISM_WASM_THREADS env is passed.</Description>
         <Tags>extism, wasm, plugin</Tags>
         <PackageLicenseExpression>BSD-3-Clause</PackageLicenseExpression>
     </PropertyGroup>

runtime/Cargo.toml

@@ -9,9 +9,10 @@ repository = "https://github.com/extism/extism"
 description = "Extism runtime component"
 
 [dependencies]
-wasmtime = "8.0.0"
-wasmtime-wasi = "8.0.0"
-wasmtime-wasi-nn = {version = "8.0.0", optional=true}
+wasmtime = "9.0.2"
+wasmtime-wasi = "9.0.2"
+wasmtime-wasi-nn = {version = "9.0.2", optional=true}
+wasmtime-wasi-threads = "9.0.2"
 anyhow = "1"
 serde = {version = "1", features = ["derive"]}
 serde_json = "1"
@@ -26,6 +27,7 @@ extism-manifest = { version = "0.3.0", path = "../manifest" }
 pretty-hex = { version = "0.3" }
 uuid = { version = "1", features = ["v4"] }
 libc = "0.2"
+rand = "0.8.5"
 
 [features]
 default = ["http", "register-http", "register-filesystem"]

runtime/src/internal.rs

@@ -1,14 +1,17 @@
 use std::collections::BTreeMap;
+use std::ptr::{null, null_mut};
+use std::sync::{Arc, Mutex, RwLock};
 
 use crate::*;
 
 /// Internal stores data that is available to the caller in PDK functions
+#[derive(Clone)]
 pub struct Internal {
     /// Call input length
     pub input_length: usize,
 
     /// Pointer to call input
-    pub input: *const u8,
+    pub input: Arc<RwLock<Option<Vec<u8>>>>,
 
     /// Memory offset that points to the output
     pub output_offset: usize,
@@ -29,7 +32,7 @@ pub struct Internal {
     pub vars: BTreeMap<String, Vec<u8>>,
 
     /// A pointer to the plugin memory, this should mostly be used from the PDK
-    pub memory: *mut PluginMemory,
+    pub memory: Arc<Mutex<Option<PluginMemory>>>,
 }
 
 /// InternalExt provides a unified way of acessing `memory`, `store` and `internal` values
@@ -65,6 +68,18 @@ pub struct Wasi {
     #[cfg(not(feature = "nn"))]
     pub nn: (),
 }
+impl Clone for Wasi {
+    fn clone(&self) -> Self {
+        Self {
+            ctx: self.ctx.clone(),
+            // I'm not sure how to handle cloning here correctly. So we will just create the new instance.
+            #[cfg(feature = "nn")]
+            nn: wasmtime_wasi_nn::WasiNnCtx::new().unwrap(),
+            #[cfg(not(feature = "nn"))]
+            nn: (),
+        }
+    }
+}
 
 impl Internal {
     pub(crate) fn new(manifest: &Manifest, wasi: bool) -> Result<Self, Error> {
@@ -101,9 +116,9 @@ impl Internal {
             input_length: 0,
             output_offset: 0,
             output_length: 0,
-            input: std::ptr::null(),
+            input: Arc::new(RwLock::from(None)),
             wasi,
-            memory: std::ptr::null_mut(),
+            memory: Arc::new(Mutex::new(None)),
             http_status: 0,
             last_error: std::cell::RefCell::new(None),
             vars: BTreeMap::new(),
@@ -123,10 +138,12 @@ impl Internal {
 
 impl InternalExt for Internal {
     fn memory(&self) -> &PluginMemory {
-        unsafe { &*self.memory }
+        let lock = self.memory.lock().unwrap();
+        lock.as_ref().unwrap()
     }
 
     fn memory_mut(&mut self) -> &mut PluginMemory {
-        unsafe { &mut *self.memory }
+        let mut lock = self.memory.lock().unwrap();
+        lock.as_mut().unwrap()
     }
 }

runtime/src/pdk.rs

@@ -40,11 +40,14 @@ pub(crate) fn input_load_u8(
     output: &mut [Val],
 ) -> Result<(), Error> {
     let data: &Internal = caller.data();
-    if data.input.is_null() {
-        return Ok(());
+    return match &*data.input.read().unwrap() {
+        None => Ok(()),
+        Some(inputData) => {
+            let val = inputData.get(input[0].unwrap_i64() as usize).unwrap_or(&0);
+            output[0] = Val::I32(*val as i32);
+            Ok(())
+        }
     }
-    output[0] = unsafe { Val::I32(*data.input.add(input[0].unwrap_i64() as usize) as i32) };
-    Ok(())
 }
 
 /// Load an unsigned 64 bit integer from input
@@ -56,14 +59,16 @@ pub(crate) fn input_load_u64(
     output: &mut [Val],
 ) -> Result<(), Error> {
     let data: &Internal = caller.data();
-    if data.input.is_null() {
-        return Ok(());
+    return match &*data.input.read().unwrap() {
+        None => Ok(()),
+        Some(inputData) => {
+            let offs = args!(input, 0, i64) as usize;
+            let slice = &inputData[offs..offs+8];
+            let byte = u64::from_ne_bytes(slice.try_into().unwrap());
+            output[0] = Val::I64(byte as i64);
+            Ok(())
+        }
     }
-    let offs = args!(input, 0, i64) as usize;
-    let slice = unsafe { std::slice::from_raw_parts(data.input.add(offs), 8) };
-    let byte = u64::from_ne_bytes(slice.try_into().unwrap());
-    output[0] = Val::I64(byte as i64);
-    Ok(())
 }
 
 /// Store a byte in memory

runtime/src/plugin.rs

@@ -1,4 +1,6 @@
 use std::collections::BTreeMap;
+use std::sync::{Arc, Mutex, RwLock};
+use wasmtime_wasi_threads::WasiThreadsCtx;
 
 use crate::*;
 
@@ -62,12 +64,18 @@ impl Plugin {
         imports: impl IntoIterator<Item = &'a Function>,
         with_wasi: bool,
     ) -> Result<Plugin, Error> {
+        // Additional configuration (mostly experimental) coming from env variables.
+        let with_wasm_threads = std::env::var("EXTISM_WASM_THREADS").is_ok();
+        let with_debug = std::env::var("EXTISM_DEBUG").is_ok();
+        let with_wasi_threads = std::env::var("EXTISM_WASI_THREADS").is_ok();
+
         // Create a new engine, if the `EXITSM_DEBUG` environment variable is set
         // then we enable debug info
         let engine = Engine::new(
             Config::new()
+                .wasm_threads(with_wasm_threads)
                 .epoch_interruption(true)
-                .debug_info(std::env::var("EXTISM_DEBUG").is_ok())
+                .debug_info(with_debug)
                 .profiler(profiling_strategy()),
         )?;
         let mut imports = imports.into_iter();
@@ -102,6 +110,12 @@ impl Plugin {
             (entry.0.as_str(), entry.1)
         });
 
+        if with_wasi && with_wasi_threads {
+            wasmtime_wasi_threads::add_to_linker(&mut linker, &mut store, &main,|x: &mut Internal| {
+                WasiThreadsCtx::new(main.clone(), Arc::new(linker.clone())).as_mut().unwrap()
+            })?;
+        }
+
         // Define PDK functions
         macro_rules! define_funcs {
             ($m:expr, { $($name:ident($($args:expr),*) $(-> $($r:expr),*)?);* $(;)?}) => {
@@ -185,7 +199,7 @@ impl Plugin {
         };
 
         // Make sure `Internal::memory` is initialized
-        plugin.internal_mut().memory = plugin.memory.get();
+        plugin.internal_mut().memory = Arc::new(Mutex::new(Some(unsafe { plugin.memory.into_inner() })));
 
         // Then detect runtime before returning the new plugin
         plugin.detect_runtime();
@@ -211,9 +225,10 @@ impl Plugin {
         }
         let ptr = self.memory.get();
         let internal = self.internal_mut();
-        internal.input = input;
         internal.input_length = len;
-        internal.memory = ptr
+        // Shit that is going here is required to ensure that input and memory is safely shared across WASI threads.
+        internal.input = Arc::new(RwLock::new(Some(unsafe { std::slice::from_raw_parts(input, len).to_vec() })));
+        internal.memory = Arc::new(Mutex::new(Some(*unsafe { Box::from_raw(ptr) })))
     }
 
     /// Dump memory using trace! logging
@@ -414,4 +429,4 @@ impl Plugin {
 pub(crate) enum Runtime {
     Haskell { init: Func, cleanup: Func },
     Wasi { init: Func, cleanup: Option<Func> },
-}
+}