diff --git a/ext/src/ruby_api/errors.rs b/ext/src/ruby_api/errors.rs index 52bcc655..f5a537f8 100644 --- a/ext/src/ruby_api/errors.rs +++ b/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 = + Lazy::new(|_| root().const_get("ModuleDisposedError").unwrap()); + let ruby = Ruby::get().unwrap(); + ruby.get_inner(&ERR) +} + +pub fn module_disposed_err() -> Result { + 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 = Lazy::new(|_| root().const_get("WasiExit").unwrap()); diff --git a/ext/src/ruby_api/instance.rs b/ext/src/ruby_api/instance.rs index 39748b58..a11ba1d4 100644 --- a/ext/src/ruby_api/instance.rs +++ b/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, &Module), (Option,), (), (), (), ()>(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))?; diff --git a/ext/src/ruby_api/linker.rs b/ext/src/ruby_api/linker.rs index 9a2c5a2e..a043e885 100644 --- a/ext/src/ruby_api/linker.rs +++ b/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)); diff --git a/ext/src/ruby_api/module.rs b/ext/src/ruby_api/module.rs index 89e30336..4087ad51 100644 --- a/ext/src/ruby_api/module.rs +++ b/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>, _track_memory_usage: ManuallyTracked<()>, } // Needed for ManuallyTracked unsafe impl Send for Module {} +unsafe impl Sync for Module {} impl Module { /// @yard @@ -95,7 +97,9 @@ impl Module { /// @return [String] /// @see .deserialize pub fn serialize(&self) -> Result { - let module = self.get(); + let Some(module) = self.get() else { + return module_disposed_err(); + }; let bytes = module.serialize(); bytes @@ -103,8 +107,25 @@ impl Module { .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 { + 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 { + Ok(self.get().is_none()) + } + + pub fn get(&self) -> Option<&ModuleImpl> { + unsafe { &*self.inner.get() }.as_ref() } } @@ -113,7 +134,7 @@ impl From 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(()) } diff --git a/lib/wasmtime/error.rb b/lib/wasmtime/error.rb index 377da9af..b50a2052 100644 --- a/lib/wasmtime/error.rb +++ b/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 diff --git a/spec/unit/module_spec.rb b/spec/unit/module_spec.rb index 57068e1d..a567d0d8 100644 --- a/spec/unit/module_spec.rb +++ b/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