fixup! retain error

Author: william-stacken

ext/src/ruby_api/store.rs

@@ -1,5 +1,6 @@
 use super::errors::wasi_exit_error;
 use super::{caller::Caller, engine::Engine, root, trap::Trap};
+use crate::ruby_api::wasi_config::WasiRetainedData;
 use crate::{define_rb_intern, error, WasiConfig};
 use magnus::value::ReprValue;
 use magnus::value::StaticSymbol;
@@ -36,6 +37,7 @@ pub struct StoreData {
     wasi_p1: Option<WasiP1Ctx>,
     wasi: Option<WasiCtx>,
     refs: Vec<Value>,
+    wasi_retained_data: Option<WasiRetainedData>,
     last_error: Option<Error>,
     store_limits: TrackingResourceLimiter,
     resource_table: ResourceTable,
@@ -73,26 +75,19 @@ impl StoreData {
     }
 
     pub fn check_socket_errors(&mut self) {
-        use crate::ruby_api::wasi_config::WasiRetainedData;
-
-        // Check all retained values for WasiRetainedData and extract error storages
-        for value in &self.refs {
-            let ruby = Ruby::get().unwrap();
-
-            // Try to convert to WasiRetainedData
-            if let Ok(retained) = Obj::<WasiRetainedData>::try_convert(*value) {
-                if let Some(storage) = retained.error_storage() {
-                    if let Ok(mut guard) = storage.lock() {
-                        if let Some((class_name, error_msg)) = guard.take() {
-                            // Look up the exception class by name and reconstruct the error
-                            let exception_class = ruby
-                                .class_object()
-                                .const_get::<_, ExceptionClass>(class_name)
-                                .unwrap_or_else(|_| ruby.exception_runtime_error());
-
-                            self.last_error = Some(Error::new(exception_class, error_msg));
-                            return;
-                        }
+        // Check the wasi_retained_data field for stored socket errors
+        if let Some(ref retained_data) = self.wasi_retained_data {
+            if let Some(storage) = retained_data.error_storage() {
+                if let Ok(mut guard) = storage.lock() {
+                    if let Some((class_name, error_msg)) = guard.take() {
+                        let ruby = Ruby::get().unwrap();
+                        // Look up the exception class by name and reconstruct the error
+                        let exception_class = ruby
+                            .class_object()
+                            .const_get::<_, ExceptionClass>(class_name)
+                            .unwrap_or_else(|_| ruby.exception_runtime_error());
+
+                        self.last_error = Some(Error::new(exception_class, error_msg));
                     }
                 }
             }
@@ -108,6 +103,10 @@ impl StoreData {
             }
         }
 
+        if let Some(ref retained_data) = self.wasi_retained_data {
+            retained_data.mark(marker);
+        }
+
         for value in self.refs.iter() {
             marker.mark_movable(*value);
         }
@@ -205,12 +204,9 @@ impl Store {
             .transpose()?
             .unzip();
 
-        // Collect any Values that need to be retained for GC
-        let refs: Vec<Value> = [wasi_retained, wasi_p1_retained]
-            .into_iter()
-            .flatten()
-            .flatten()
-            .collect();
+        // Store WasiRetainedData directly (prefer wasi over wasi_p1 if both present)
+        let wasi_retained_data: Option<WasiRetainedData> =
+            wasi_retained.flatten().or(wasi_p1_retained.flatten());
 
         let limiter = match kw.optional.2 {
             None => StoreLimitsBuilder::new(),
@@ -224,7 +220,8 @@ impl Store {
             user_data,
             wasi_p1,
             wasi,
-            refs,
+            refs: Default::default(),
+            wasi_retained_data,
             last_error: Default::default(),
             store_limits: limiter,
             resource_table: Default::default(),

ext/src/ruby_api/wasi_config.rs

@@ -35,22 +35,17 @@ type SocketErrorStorage = Arc<Mutex<Option<(String, String)>>>;
 
 /// Container for data that needs to be retained by the Store for WASI functionality.
 /// This includes Ruby procs (for GC marking) and error storages (for error propagation).
-#[derive(TypedData)]
-#[magnus(class = "Wasmtime::WasiRetainedData", mark, free_immediately)]
 pub struct WasiRetainedData {
     proc: Option<Opaque<Proc>>,
     error_storage: Option<SocketErrorStorage>,
 }
 
-impl DataTypeFunctions for WasiRetainedData {
-    fn mark(&self, marker: &Marker) {
+impl WasiRetainedData {
+    pub fn mark(&self, marker: &Marker) {
         if let Some(proc) = self.proc {
             marker.mark(proc);
         }
     }
-}
-
-impl WasiRetainedData {
     pub fn new(proc: Option<Proc>, error_storage: Option<SocketErrorStorage>) -> Self {
         Self {
             proc: proc.map(|p| p.into()),
@@ -505,19 +500,19 @@ impl WasiConfig {
         rb_self
     }
 
-    pub fn build_p1(&self, ruby: &Ruby) -> Result<(WasiP1Ctx, Option<Value>), Error> {
+    pub fn build_p1(&self, ruby: &Ruby) -> Result<(WasiP1Ctx, Option<WasiRetainedData>), Error> {
         let (mut builder, retained_data) = self.build_impl(ruby)?;
         let ctx = builder.build_p1();
         Ok((ctx, retained_data))
     }
 
-    pub fn build(&self, ruby: &Ruby) -> Result<(WasiCtx, Option<Value>), Error> {
+    pub fn build(&self, ruby: &Ruby) -> Result<(WasiCtx, Option<WasiRetainedData>), Error> {
         let (mut builder, retained_data) = self.build_impl(ruby)?;
         let ctx = builder.build();
         Ok((ctx, retained_data))
     }
 
-    fn build_impl(&self, ruby: &Ruby) -> Result<(WasiCtxBuilder, Option<Value>), Error> {
+    fn build_impl(&self, ruby: &Ruby) -> Result<(WasiCtxBuilder, Option<WasiRetainedData>), Error> {
         let mut builder = WasiCtxBuilder::new();
         let inner = self.inner.borrow();
         let mut proc_to_retain = None;
@@ -631,15 +626,14 @@ impl WasiConfig {
                 .map_err(|e| error!("{}", e))?;
         }
 
-        // Wrap retained data in a single Ruby object if we have anything to retain
-        let retained_value = if proc_to_retain.is_some() || error_storage_to_retain.is_some() {
-            let retained_data = WasiRetainedData::new(proc_to_retain, error_storage_to_retain);
-            Some(ruby.obj_wrap(retained_data).as_value())
+        // Return retained data directly if we have anything to retain
+        let retained_data = if proc_to_retain.is_some() || error_storage_to_retain.is_some() {
+            Some(WasiRetainedData::new(proc_to_retain, error_storage_to_retain))
         } else {
             None
         };
 
-        Ok((builder, retained_value))
+        Ok((builder, retained_data))
     }
 
     fn check_determinism(&self) -> Result<(), Error> {
@@ -671,9 +665,6 @@ pub fn file_w(path: RString) -> Result<File, Error> {
 }
 
 pub fn init(ruby: &Ruby) -> Result<(), Error> {
-    // Register internal WasiRetainedData class (not exposed to Ruby API)
-    root().define_class("WasiRetainedData", ruby.class_object())?;
-
     let class = root().define_class("WasiConfig", ruby.class_object())?;
     class.define_singleton_method("new", function!(WasiConfig::new, 0))?;