Allow for manually disposing of Wasmtime::Module

ext/src/ruby_api/errors.rs

@@ -23,6 +23,21 @@ pub fn conversion_error() -> ExceptionClass {
     ruby.get_inner(&ERR)
 }
 
+/// Raised when attempting to use a module that has been disposed.
+pub fn module_disposed_error() -> ExceptionClass {
+    static ERR: Lazy<ExceptionClass> =
+        Lazy::new(|_| root().const_get("ModuleDisposedError").unwrap());
+    let ruby = Ruby::get().unwrap();
+    ruby.get_inner(&ERR)
+}
+
+pub fn module_disposed_err<T>() -> Result<T, Error> {
+    Err(Error::new(
+        module_disposed_error(),
+        "module has been disposed",
+    ))
+}
+
 /// Raised when a WASI program terminates early by calling +exit+.
 pub fn wasi_exit_error() -> ExceptionClass {
     static ERR: Lazy<ExceptionClass> = Lazy::new(|_| root().const_get("WasiExit").unwrap());

ext/src/ruby_api/instance.rs

@@ -1,5 +1,6 @@
 use super::{
     convert::{ToExtern, WrapWasmtimeType},
+    errors::module_disposed_err,
     func::Func,
     module::Module,
     root,
@@ -42,6 +43,7 @@ impl Instance {
         let args =
             scan_args::scan_args::<(Obj<Store>, &Module), (Option<Value>,), (), (), (), ()>(args)?;
         let (wrapped_store, module) = args.required;
+
         let mut context = wrapped_store.context_mut();
         let imports = args
             .optional
@@ -62,7 +64,9 @@ impl Instance {
             None => vec![],
         };
 
-        let module = module.get();
+        let Some(module) = module.get() else {
+            return module_disposed_err();
+        };
         let inner = InstanceImpl::new(context, module, &imports)
             .map_err(|e| StoreContextValue::from(wrapped_store).handle_wasm_error(e))?;
 

ext/src/ruby_api/linker.rs

@@ -2,6 +2,7 @@ use super::{
     convert::WrapWasmtimeType,
     convert::{ToExtern, ToValTypeVec},
     engine::Engine,
+    errors::module_disposed_err,
     externals::Extern,
     func::{self, Func},
     instance::Instance,
@@ -11,9 +12,9 @@ use super::{
 };
 use crate::{define_rb_intern, err, error};
 use magnus::{
-    block::Proc, class, function, gc::Marker, method, prelude::*, scan_args, scan_args::scan_args,
-    typed_data::Obj, DataTypeFunctions, Error, Object, RArray, RHash, RString, Ruby, TypedData,
-    Value,
+    block::Proc, class, function, gc::Marker, method, module, prelude::*, scan_args,
+    scan_args::scan_args, typed_data::Obj, DataTypeFunctions, Error, Object, RArray, RHash,
+    RString, Ruby, TypedData, Value,
 };
 use std::cell::RefCell;
 use wasmtime::Linker as LinkerImpl;
@@ -86,9 +87,13 @@ impl Linker {
     /// @param mod [Module]
     /// @return [void]
     pub fn define_unknown_imports_as_traps(&self, module: &Module) -> Result<(), Error> {
+        let Some(module) = module.get() else {
+            return module_disposed_err();
+        };
+
         self.inner
             .borrow_mut()
-            .define_unknown_imports_as_traps(module.get())
+            .define_unknown_imports_as_traps(module)
             .map_err(|e| error!("{}", e))
     }
 
@@ -218,9 +223,13 @@ impl Linker {
     /// @param mod [Module]
     /// @return [void]
     pub fn module(&self, store: &Store, name: RString, module: &Module) -> Result<(), Error> {
+        let Some(module) = module.get() else {
+            return module_disposed_err();
+        };
+
         self.inner
             .borrow_mut()
-            .module(store.context_mut(), unsafe { name.as_str()? }, module.get())
+            .module(store.context_mut(), unsafe { name.as_str()? }, module)
             .map(|_| ())
             .map_err(|e| error!("{}", e))
     }
@@ -285,9 +294,13 @@ impl Linker {
             );
         }
 
+        let Some(module) = module.get() else {
+            return module_disposed_err();
+        };
+
         self.inner
             .borrow_mut()
-            .instantiate(store.context_mut(), module.get())
+            .instantiate(store.context_mut(), module)
             .map_err(|e| StoreContextValue::from(store).handle_wasm_error(e))
             .map(|instance| {
                 self.refs.borrow().iter().for_each(|val| store.retain(*val));

ext/src/ruby_api/module.rs

@@ -1,10 +1,12 @@
 use std::{
+    borrow::Borrow,
+    cell::{Cell, RefCell, UnsafeCell},
     mem::{transmute, MaybeUninit},
     ops::Deref,
     os::raw::c_void,
 };
 
-use super::{engine::Engine, root};
+use super::{engine::Engine, errors::module_disposed_err, root};
 use crate::{
     error,
     helpers::{nogvl, Tmplock},
@@ -18,15 +20,15 @@ use wasmtime::Module as ModuleImpl;
 /// @yard
 /// Represents a WebAssembly module.
 /// @see https://docs.rs/wasmtime/latest/wasmtime/struct.Module.html Wasmtime's Rust doc
-#[derive(Clone)]
 #[magnus::wrap(class = "Wasmtime::Module", size, free_immediately, frozen_shareable)]
 pub struct Module {
-    inner: ModuleImpl,
+    inner: UnsafeCell<Option<ModuleImpl>>,
     _track_memory_usage: ManuallyTracked<()>,
 }
 
 // Needed for ManuallyTracked
 unsafe impl Send for Module {}
+unsafe impl Sync for Module {}
 
 impl Module {
     /// @yard
@@ -95,16 +97,35 @@ impl Module {
     /// @return [String]
     /// @see .deserialize
     pub fn serialize(&self) -> Result<RString, Error> {
-        let module = self.get();
+        let Some(module) = self.get() else {
+            return module_disposed_err();
+        };
         let bytes = module.serialize();
 
         bytes
             .map(|bytes| RString::from_slice(&bytes))
             .map_err(|e| error!("{:?}", e))
     }
 
-    pub fn get(&self) -> &ModuleImpl {
-        &self.inner
+    /// @yard
+    /// Disposes the module, freeing the underlying resources without waiting
+    /// for the GC. The module is rendered unusable after this call.
+    /// @return [Boolean] Whether the module was disposed or not.
+    pub fn dispose(&self) -> Result<bool, Error> {
+        let borrowed = unsafe { &mut *self.inner.get() };
+
+        Ok(borrowed.take().is_some())
+    }
+
+    /// @yard
+    /// Checks whether the module has been disposed manually.
+    /// @return [Boolean] Whether the module has been manually disposed of.
+    pub fn is_disposed(&self) -> Result<bool, Error> {
+        Ok(self.get().is_none())
+    }
+
+    pub fn get(&self) -> Option<&ModuleImpl> {
+        unsafe { &*self.inner.get() }.as_ref()
     }
 }
 
@@ -113,7 +134,7 @@ impl From<ModuleImpl> for Module {
         let size = inner.image_range().len();
 
         Self {
-            inner,
+            inner: UnsafeCell::new(Some(inner)),
             _track_memory_usage: ManuallyTracked::new(size),
         }
     }
@@ -127,6 +148,8 @@ pub fn init() -> Result<(), Error> {
     class.define_singleton_method("deserialize", function!(Module::deserialize, 2))?;
     class.define_singleton_method("deserialize_file", function!(Module::deserialize_file, 2))?;
     class.define_method("serialize", method!(Module::serialize, 0))?;
+    class.define_method("dispose!", method!(Module::dispose, 0))?;
+    class.define_method("disposed?", method!(Module::is_disposed, 0))?;
 
     Ok(())
 }

lib/wasmtime/error.rb

@@ -45,4 +45,7 @@ def message
       "WASI exit with code #{code}"
     end
   end
+
+  # Raised when attempting to use a disposed Wasmtime module.
+  class ModuleDisposedError < Error; end
 end

spec/unit/module_spec.rb

@@ -86,5 +86,22 @@ def create_tmpfile(content)
         expect(mod).to be_a(Wasmtime::Module)
       end
     end
+
+    describe "#dispose" do
+      it "renders the module unusable" do
+        mod = Module.new(engine, wat)
+        mod.dispose!
+        expect { mod.serialize }.to raise_error(Wasmtime::Error)
+      end
+
+      it "can be called multiple times" do
+        mod = Module.new(engine, wat)
+
+        expect(mod.disposed?).to be(false)
+        expect(mod.dispose!).to be(true)
+        expect(mod.disposed?).to be(true)
+        expect(mod.dispose!).to be(false)
+      end
+    end
   end
 end