Change API to work with wasmtime-wasi (#467)

* Change API to work with wasmtime-wasi

* Use 33.0.0 for Rust crate version

Author: jeffcharles

Cargo.lock

@@ -9,12 +9,11 @@
 mod1 = Wasmtime::Module.from_file(engine, "examples/linking1.wat")
 mod2 = Wasmtime::Module.from_file(engine, "examples/linking2.wat")
 
-wasi_ctx_builder = Wasmtime::WasiCtxBuilder.new
+wasi_config = Wasmtime::WasiConfig.new
   .inherit_stdin
   .inherit_stdout
-  .build
 
-store = Wasmtime::Store.new(engine, wasi_ctx: wasi_ctx_builder)
+store = Wasmtime::Store.new(engine, wasi_config: wasi_config)
 
 # Instantiate `mod2` which only uses WASI, then register
 # that instance with the linker so `mod1` can use it.

examples/linking.rb

@@ -5,14 +5,13 @@
 
 linker = Wasmtime::Linker.new(engine, wasi: true)
 
-wasi_ctx = Wasmtime::WasiCtxBuilder.new
+wasi_config = Wasmtime::WasiConfig.new
   .set_stdin_string("hi!")
   .inherit_stdout
   .inherit_stderr
   .set_argv(ARGV)
   .set_env(ENV)
-  .build
-store = Wasmtime::Store.new(engine, wasi_ctx: wasi_ctx)
+store = Wasmtime::Store.new(engine, wasi_config: wasi_config)
 
 instance = linker.instantiate(store, mod)
 instance.invoke("_start")

examples/wasi.rb

@@ -1,6 +1,6 @@
 [package]
 name = "wasmtime-rb"
-version = "9.0.4"
+version = "33.0.0"
 edition = "2021"
 authors = ["The Wasmtime Project Developers"]
 license = "Apache-2.0"
@@ -17,14 +17,15 @@ all-arch = ["wasmtime/all-arch"]
 winch = ["wasmtime/winch"]
 
 [dependencies]
+async-trait = "*" # Needed for `OutputLimitedBuffer`. Use wasmtime's version.
+bytes = "*" # Needed for `OutputLimitedBuffer`. Use wasmtime's version.
 lazy_static = "1.5.0"
 magnus = { version = "0.7", features = ["rb-sys"] }
 rb-sys = { version = "*", default-features = false, features = [
   "stable-api-compiled-fallback",
 ] }
 wasmtime = { version = "=33.0.0", features = ["memory-protection-keys"] }
 wasmtime-wasi = "=33.0.0"
-wasi-common = "=33.0.0"
 cap-std = "3.4.0"
 wat = "1.227.1"
 tokio = { version = "1.40.0", features = [
@@ -38,7 +39,7 @@ async-timer = { version = "1.0.0-beta.15", features = [
 ], optional = true }
 static_assertions = "1.1.0"
 wasmtime-environ = "=33.0.0"
-deterministic-wasi-ctx = { version = "=1.2.1", features = ["wasi-common"] }
+deterministic-wasi-ctx = { version = "=1.2.1" }
 
 [build-dependencies]
 rb-sys-env = "0.2.2"

ext/Cargo.toml

@@ -1,28 +1,88 @@
+use bytes::Bytes;
 use magnus::{
     value::{InnerValue, Opaque, ReprValue},
     RString, Ruby,
 };
-use std::io;
-use std::io::ErrorKind;
+use std::io::Write;
+use std::sync::{Arc, Mutex};
+use wasmtime_wasi::p2::{OutputStream, Pollable, StdoutStream, StreamError, StreamResult};
 
 /// A buffer that limits the number of bytes that can be written to it.
 /// If the buffer is full, it will truncate the data.
-/// Is used in the buffer implementations of stdout and stderr in `WasiCtx` and `WasiCtxBuilder`.
+/// Is used in the buffer implementations of stdout and stderr in `WasiP1Ctx` and `WasiCtxBuilder`.
 pub struct OutputLimitedBuffer {
+    inner: Arc<Mutex<OutputLimitedBufferInner>>,
+}
+
+impl OutputLimitedBuffer {
+    /// Creates a new [OutputLimitedBuffer] with the given underlying buffer
+    /// and capacity.
+    pub fn new(buffer: Opaque<RString>, capacity: usize) -> Self {
+        Self {
+            inner: Arc::new(Mutex::new(OutputLimitedBufferInner::new(buffer, capacity))),
+        }
+    }
+}
+
+impl Clone for OutputLimitedBuffer {
+    fn clone(&self) -> Self {
+        Self {
+            inner: Arc::clone(&self.inner),
+        }
+    }
+}
+
+impl StdoutStream for OutputLimitedBuffer {
+    fn stream(&self) -> Box<dyn OutputStream> {
+        let cloned = self.clone();
+        Box::new(cloned)
+    }
+
+    fn isatty(&self) -> bool {
+        false
+    }
+}
+
+#[async_trait::async_trait]
+impl Pollable for OutputLimitedBuffer {
+    async fn ready(&mut self) {}
+}
+
+impl OutputStream for OutputLimitedBuffer {
+    fn write(&mut self, bytes: Bytes) -> StreamResult<()> {
+        let mut stream = self.inner.lock().expect("Should be only writer");
+        stream.write(&bytes)
+    }
+
+    fn flush(&mut self) -> StreamResult<()> {
+        Ok(())
+    }
+
+    fn check_write(&mut self) -> StreamResult<usize> {
+        let mut stream = self.inner.lock().expect("Should be only writer");
+        stream.check_write()
+    }
+}
+
+struct OutputLimitedBufferInner {
     buffer: Opaque<RString>,
     /// The maximum number of bytes that can be written to the output stream buffer.
     capacity: usize,
 }
 
-impl OutputLimitedBuffer {
+impl OutputLimitedBufferInner {
     #[must_use]
     pub fn new(buffer: Opaque<RString>, capacity: usize) -> Self {
         Self { buffer, capacity }
     }
 }
 
-impl io::Write for OutputLimitedBuffer {
-    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+impl OutputLimitedBufferInner {
+    fn check_write(&mut self) -> StreamResult<usize> {
+        Ok(usize::MAX)
+    }
+
+    fn write(&mut self, buf: &[u8]) -> StreamResult<()> {
         // Append a buffer to the string and truncate when hitting the capacity.
         // We return the input buffer size regardless of whether we truncated or not to avoid a panic.
         let ruby = Ruby::get().unwrap();
@@ -32,39 +92,28 @@ impl io::Write for OutputLimitedBuffer {
         // Handling frozen case here is necessary because magnus does not check if a string is frozen before writing to it.
         let is_frozen = inner_buffer.as_value().is_frozen();
         if is_frozen {
-            return Err(io::Error::new(
-                ErrorKind::WriteZero,
-                "Cannot write to a frozen buffer.",
-            ));
+            return Err(StreamError::trap("Cannot write to a frozen buffer."));
         }
 
         if buf.is_empty() {
-            return Ok(0);
+            return Ok(());
         }
 
         if inner_buffer
             .len()
             .checked_add(buf.len())
             .is_some_and(|val| val < self.capacity)
         {
-            let amount_written = inner_buffer.write(buf)?;
-            if amount_written < buf.len() {
-                return Ok(amount_written);
-            }
+            inner_buffer
+                .write(buf)
+                .map_err(|e| StreamError::trap(&e.to_string()))?;
         } else {
             let portion = self.capacity - inner_buffer.len();
-            let amount_written = inner_buffer.write(&buf[0..portion])?;
-            if amount_written < portion {
-                return Ok(amount_written);
-            }
+            inner_buffer
+                .write(&buf[0..portion])
+                .map_err(|e| StreamError::trap(&e.to_string()))?;
         };
 
-        Ok(buf.len())
-    }
-
-    fn flush(&mut self) -> io::Result<()> {
-        let ruby = Ruby::get().unwrap();
-
-        self.buffer.get_inner_with(&ruby).flush()
+        Ok(())
     }
 }

ext/src/helpers/output_limited_buffer.rs

@@ -54,11 +54,8 @@ impl Linker {
 
         let mut inner: LinkerImpl<StoreData> = LinkerImpl::new(engine.get());
         if wasi {
-            wasi_common::sync::snapshots::preview_1::add_wasi_snapshot_preview1_to_linker(
-                &mut inner,
-                |s| s.wasi_ctx_mut(),
-            )
-            .map_err(|e| error!("{}", e))?
+            wasmtime_wasi::preview1::add_to_linker_sync(&mut inner, |s| s.wasi_ctx_mut())
+                .map_err(|e| error!("{}", e))?
         }
         Ok(Self {
             inner: RefCell::new(inner),
@@ -288,8 +285,8 @@ impl Linker {
                 "Store is missing WASI configuration.\n\n\
                 When using `wasi: true`, the Store given to\n\
                 `Linker#instantiate` must have a WASI configuration.\n\
-                To fix this, provide the `wasi_ctx` when creating the Store:\n\
-                    Wasmtime::Store.new(engine, wasi_ctx: WasiCtxBuilder.new)"
+                To fix this, provide the `wasi_config` when creating the Store:\n\
+                    Wasmtime::Store.new(engine, wasi_config: WasiConfig.new)"
             );
         }
 
@@ -316,6 +313,15 @@ impl Linker {
             .map(|func| Func::from_inner(store.into(), func))
             .map_err(|e| error!("{}", e))
     }
+
+    /// @yard
+    /// Replaces the `poll_oneoff` and `sched_yield` function implementations
+    /// with deterministic ones.
+    /// @return [void]
+    pub fn use_deterministic_scheduling_functions(&self) -> Result<(), Error> {
+        let mut inner = self.inner.borrow_mut();
+        deterministic_wasi_ctx::replace_scheduling_functions(&mut inner).map_err(|e| error!("{e}"))
+    }
 }
 
 pub fn init() -> Result<(), Error> {
@@ -339,6 +345,10 @@ pub fn init() -> Result<(), Error> {
     class.define_method("alias_module", method!(Linker::alias_module, 2))?;
     class.define_method("instantiate", method!(Linker::instantiate, 2))?;
     class.define_method("get_default", method!(Linker::get_default, 2))?;
+    class.define_method(
+        "use_deterministic_scheduling_functions",
+        method!(Linker::use_deterministic_scheduling_functions, 0),
+    )?;
 
     Ok(())
 }

ext/src/ruby_api/linker.rs

@@ -27,8 +27,7 @@ mod pooling_allocation_config;
 mod store;
 mod table;
 mod trap;
-mod wasi_ctx;
-mod wasi_ctx_builder;
+mod wasi_config;
 
 pub use caller::Caller;
 pub use engine::Engine;
@@ -41,8 +40,7 @@ pub use params::Params;
 pub use pooling_allocation_config::PoolingAllocationConfig;
 pub use store::Store;
 pub use trap::Trap;
-pub use wasi_ctx::WasiCtx;
-pub use wasi_ctx_builder::WasiCtxBuilder;
+pub use wasi_config::WasiConfig;
 
 /// The "Wasmtime" Ruby module.
 pub fn root() -> RModule {
@@ -85,10 +83,9 @@ pub fn init(ruby: &Ruby) -> Result<(), Error> {
     memory::init(ruby)?;
     linker::init()?;
     externals::init()?;
-    wasi_ctx_builder::init()?;
+    wasi_config::init()?;
     table::init()?;
     global::init()?;
-    wasi_ctx::init()?;
     pooling_allocation_config::init()?;
     component::init(ruby)?;