From 270b5fe2db4c1c93440f2fe77084680dc1353574 Mon Sep 17 00:00:00 2001
From: Pete Matsyburka <pete.matsy@gmail.com>
Date: Sat, 4 Jul 2026 21:53:14 +0300
Subject: [PATCH] Add Store#close

---
 ext/src/ruby_api/component/func.rs         |   4 +-
 ext/src/ruby_api/component/instance.rs     |  11 ++-
 ext/src/ruby_api/component/linker.rs       |   4 +-
 ext/src/ruby_api/component/wasi_command.rs |   6 +-
 ext/src/ruby_api/func.rs                   |   2 +-
 ext/src/ruby_api/global.rs                 |   2 +-
 ext/src/ruby_api/instance.rs               |   8 +-
 ext/src/ruby_api/linker.rs                 |  18 ++--
 ext/src/ruby_api/memory.rs                 |   2 +-
 ext/src/ruby_api/store.rs                  | 108 +++++++++++++++------
 ext/src/ruby_api/table.rs                  |   2 +-
 spec/unit/store_spec.rb                    |  46 +++++++++
 12 files changed, 156 insertions(+), 57 deletions(-)

diff --git a/ext/src/ruby_api/component/func.rs b/ext/src/ruby_api/component/func.rs
index 9e115b68..f984149f 100644
--- a/ext/src/ruby_api/component/func.rs
+++ b/ext/src/ruby_api/component/func.rs
@@ -97,12 +97,12 @@ impl Func {
         args: &[Value],
     ) -> Result<Value, Error> {
         let store_context_value = StoreContextValue::from(store);
-        let func_ty = func.ty(store.context_mut());
+        let func_ty = func.ty(store.context_mut()?);
         let results_ty = func_ty.results();
         let mut results = vec![wasmtime::component::Val::Bool(false); results_ty.len()];
         let params = convert_params(ruby, &store_context_value, func_ty.params(), args)?;
 
-        func.call(store.context_mut(), &params, &mut results)
+        func.call(store.context_mut()?, &params, &mut results)
             .map_err(|e| store_context_value.handle_wasm_error(ruby, e))?;
 
         // Check for any errors stored during execution (e.g., from socket checks)
diff --git a/ext/src/ruby_api/component/instance.rs b/ext/src/ruby_api/component/instance.rs
index 32106012..a93bbdf0 100644
--- a/ext/src/ruby_api/component/instance.rs
+++ b/ext/src/ruby_api/component/instance.rs
@@ -54,9 +54,12 @@ impl Instance {
     /// @example Retrieve an +add+ export nested under an +adder+ instance top-level export:
     ///   instance.get_func(["adder", "add"])
     pub fn get_func(rb_self: Obj<Self>, handle: Value) -> Result<Option<Func>, Error> {
+        let Some(index) = rb_self.export_index(handle)? else {
+            return Ok(None);
+        };
         let func = rb_self
-            .export_index(handle)?
-            .and_then(|index| rb_self.inner.get_func(rb_self.store.context_mut(), index))
+            .inner
+            .get_func(rb_self.store.context_mut()?, index)
             .map(|inner| Func::from_inner(inner, rb_self, rb_self.store));
 
         Ok(func)
@@ -76,7 +79,7 @@ impl Instance {
 
         let index = if let Some(name) = RString::from_value(handle) {
             self.inner
-                .get_export(self.store.context_mut(), None, unsafe { name.as_str()? })
+                .get_export(self.store.context_mut()?, None, unsafe { name.as_str()? })
                 .map(|(_, index)| index)
         } else if let Some(names) = RArray::from_value(handle) {
             unsafe { names.as_slice() }
@@ -86,7 +89,7 @@ impl Instance {
 
                     Ok(self
                         .inner
-                        .get_export(self.store.context_mut(), index.as_ref(), unsafe {
+                        .get_export(self.store.context_mut()?, index.as_ref(), unsafe {
                             name.as_str()?
                         })
                         .map(|(_, index)| index))
diff --git a/ext/src/ruby_api/component/linker.rs b/ext/src/ruby_api/component/linker.rs
index f89616eb..dc70335c 100644
--- a/ext/src/ruby_api/component/linker.rs
+++ b/ext/src/ruby_api/component/linker.rs
@@ -137,13 +137,13 @@ impl Linker {
         store: Obj<Store>,
         component: &Component,
     ) -> Result<Instance, Error> {
-        if *rb_self.has_wasi.borrow() && !store.context().data().has_wasi_ctx() {
+        if *rb_self.has_wasi.borrow() && !store.context()?.data().has_wasi_ctx() {
             return err!("{}", errors::missing_wasi_ctx_error("linker.instantiate"));
         }
 
         let inner = rb_self.inner.borrow();
         inner
-            .instantiate(store.context_mut(), component.get())
+            .instantiate(store.context_mut()?, component.get())
             .map(|instance| {
                 rb_self
                     .refs
diff --git a/ext/src/ruby_api/component/wasi_command.rs b/ext/src/ruby_api/component/wasi_command.rs
index fb1b55b7..21aeac89 100644
--- a/ext/src/ruby_api/component/wasi_command.rs
+++ b/ext/src/ruby_api/component/wasi_command.rs
@@ -27,11 +27,11 @@ impl WasiCommand {
     /// @param linker [Linker]
     /// @return [WasiCommand]
     pub fn new(store: &Store, component: &Component, linker: &Linker) -> Result<Self, Error> {
-        if linker.has_wasi() && !store.context().data().has_wasi_ctx() {
+        if linker.has_wasi() && !store.context()?.data().has_wasi_ctx() {
             return err!("{}", errors::missing_wasi_ctx_error("WasiCommand.new"));
         }
         let command =
-            Command::instantiate(store.context_mut(), component.get(), &linker.inner_mut())
+            Command::instantiate(store.context_mut()?, component.get(), &linker.inner_mut())
                 .map_err(|e| error!("{e}"))?;
         Ok(Self { command })
     }
@@ -45,7 +45,7 @@ impl WasiCommand {
         rb_self
             .command
             .wasi_cli_run()
-            .call_run(store.context_mut())
+            .call_run(store.context_mut()?)
             .map_err(|err| error!("{err}"))?
             .map_err(|_| error!("Error running `run`"))?;
 
diff --git a/ext/src/ruby_api/func.rs b/ext/src/ruby_api/func.rs
index 8fd77df2..58f4c01b 100644
--- a/ext/src/ruby_api/func.rs
+++ b/ext/src/ruby_api/func.rs
@@ -136,7 +136,7 @@ impl<'a> Func<'a> {
 
         store.retain(callable.as_value());
 
-        let context = store.context_mut();
+        let context = store.context_mut()?;
         let engine = context.engine();
         let ty = wasmtime::FuncType::new(
             engine,
diff --git a/ext/src/ruby_api/global.rs b/ext/src/ruby_api/global.rs
index 00edfd17..a36b5d3b 100644
--- a/ext/src/ruby_api/global.rs
+++ b/ext/src/ruby_api/global.rs
@@ -105,7 +105,7 @@ impl<'a> Global<'a> {
         let wasm_type = value_type.to_val_type()?;
         let wasm_default = default.to_wasm_val(&store.into(), wasm_type.clone())?;
         let inner = GlobalImpl::new(
-            store.context_mut(),
+            store.context_mut()?,
             wasmtime::GlobalType::new(wasm_type, mutability),
             wasm_default,
         )
diff --git a/ext/src/ruby_api/instance.rs b/ext/src/ruby_api/instance.rs
index 46a699e6..7fe9d626 100644
--- a/ext/src/ruby_api/instance.rs
+++ b/ext/src/ruby_api/instance.rs
@@ -42,7 +42,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 mut context = wrapped_store.context_mut()?;
         let imports = args
             .optional
             .0
@@ -87,7 +87,7 @@ impl Instance {
     /// @def exports
     /// @return [Hash{String => Extern}]
     pub fn exports(ruby: &Ruby, rb_self: Obj<Self>) -> Result<RHash, Error> {
-        let mut ctx = rb_self.store.context_mut();
+        let mut ctx = rb_self.store.context_mut()?;
         let hash = ruby.hash_new();
 
         for export in rb_self.inner.exports(&mut ctx) {
@@ -111,7 +111,7 @@ impl Instance {
     pub fn export(&self, str: RString) -> Result<Option<super::externals::Extern<'_>>, Error> {
         let export = self
             .inner
-            .get_export(self.store.context_mut(), unsafe { str.as_str()? });
+            .get_export(self.store.context_mut()?, unsafe { str.as_str()? });
         let ruby = Ruby::get_with(str);
         match export {
             Some(export) => export
@@ -138,7 +138,7 @@ impl Instance {
             )
         })?)?;
 
-        let func = rb_self.get_func(rb_self.store.context_mut(), unsafe { name.as_str()? })?;
+        let func = rb_self.get_func(rb_self.store.context_mut()?, unsafe { name.as_str()? })?;
         Func::invoke(ruby, &rb_self.store.into(), &func, true, &args[1..])
     }
 
diff --git a/ext/src/ruby_api/linker.rs b/ext/src/ruby_api/linker.rs
index 3a8afa6a..5cdec806 100644
--- a/ext/src/ruby_api/linker.rs
+++ b/ext/src/ruby_api/linker.rs
@@ -100,7 +100,7 @@ impl Linker {
             .inner
             .borrow_mut()
             .define(
-                store.context(),
+                store.context()?,
                 unsafe { module.as_str()? },
                 unsafe { name.as_str()? },
                 item,
@@ -167,7 +167,7 @@ impl Linker {
         let ext = self
             .inner
             .borrow()
-            .get(store.context_mut(), unsafe { module.as_str() }?, unsafe {
+            .get(store.context_mut()?, unsafe { module.as_str() }?, unsafe {
                 name.as_str()?
             })
             .map_err(|e| error!("{}", e))?;
@@ -193,7 +193,7 @@ impl Linker {
         self.inner
             .borrow_mut()
             .instance(
-                store.context_mut(),
+                store.context_mut()?,
                 unsafe { module.as_str() }?,
                 instance.get(),
             )
@@ -212,7 +212,11 @@ impl Linker {
     pub fn module(&self, store: &Store, name: RString, module: &Module) -> Result<(), Error> {
         self.inner
             .borrow_mut()
-            .module(store.context_mut(), unsafe { name.as_str()? }, module.get())
+            .module(
+                store.context_mut()?,
+                unsafe { name.as_str()? },
+                module.get(),
+            )
             .map(|_| ())
             .map_err(|e| error!("{}", e))
     }
@@ -272,14 +276,14 @@ impl Linker {
         store: Obj<Store>,
         module: &Module,
     ) -> Result<Instance, Error> {
-        if *rb_self.has_wasi.borrow() && !store.context().data().has_wasi_p1_ctx() {
+        if *rb_self.has_wasi.borrow() && !store.context()?.data().has_wasi_p1_ctx() {
             return err!("{}", errors::missing_wasi_p1_ctx_error());
         }
 
         rb_self
             .inner
             .borrow_mut()
-            .instantiate(store.context_mut(), module.get())
+            .instantiate(store.context_mut()?, module.get())
             .map_err(|e| StoreContextValue::from(store).handle_wasm_error(ruby, e))
             .map(|instance| {
                 rb_self
@@ -300,7 +304,7 @@ impl Linker {
     pub fn get_default(&self, store: Obj<Store>, module: RString) -> Result<Func<'_>, Error> {
         self.inner
             .borrow()
-            .get_default(store.context_mut(), unsafe { module.as_str() }?)
+            .get_default(store.context_mut()?, unsafe { module.as_str() }?)
             .map(|func| Func::from_inner(store.into(), func))
             .map_err(|e| error!("{}", e))
     }
diff --git a/ext/src/ruby_api/memory.rs b/ext/src/ruby_api/memory.rs
index 52b560ac..8b1761f6 100644
--- a/ext/src/ruby_api/memory.rs
+++ b/ext/src/ruby_api/memory.rs
@@ -94,7 +94,7 @@ impl<'a> Memory<'a> {
 
         let memtype = wasmtime::MemoryType::new(min, max);
 
-        let inner = MemoryImpl::new(store.context_mut(), memtype).map_err(|e| error!("{}", e))?;
+        let inner = MemoryImpl::new(store.context_mut()?, memtype).map_err(|e| error!("{}", e))?;
 
         Ok(Self {
             store: store.into(),
diff --git a/ext/src/ruby_api/store.rs b/ext/src/ruby_api/store.rs
index 822e22c2..24978623 100644
--- a/ext/src/ruby_api/store.rs
+++ b/ext/src/ruby_api/store.rs
@@ -128,16 +128,20 @@ impl StoreData {
 #[derive(Debug, TypedData)]
 #[magnus(class = "Wasmtime::Store", size, mark, compact, free_immediately)]
 pub struct Store {
-    inner: UnsafeCell<StoreImpl<StoreData>>,
+    inner: UnsafeCell<Option<StoreImpl<StoreData>>>,
 }
 
 impl DataTypeFunctions for Store {
     fn mark(&self, marker: &Marker) {
-        self.context().data().mark(marker);
+        if let Some(inner) = unsafe { (*self.inner.get()).as_ref() } {
+            inner.data().mark(marker);
+        }
     }
 
     fn compact(&self, compactor: &Compactor) {
-        self.context_mut().data_mut().compact(compactor);
+        if let Some(inner) = unsafe { (*self.inner.get()).as_mut() } {
+            inner.data_mut().compact(compactor);
+        }
     }
 }
 
@@ -231,18 +235,18 @@ impl Store {
             resource_table: Default::default(),
         };
         let store = Self {
-            inner: UnsafeCell::new(StoreImpl::new(eng, store_data)),
+            inner: UnsafeCell::new(Some(StoreImpl::new(eng, store_data))),
         };
 
-        unsafe { &mut *store.inner.get() }.limiter(|data| &mut data.store_limits);
+        unsafe { (*store.inner.get()).as_mut().unwrap() }.limiter(|data| &mut data.store_limits);
 
         Ok(store)
     }
 
     /// @yard
     /// @return [Object] The passed in value in {.new}
-    pub fn data(&self) -> Value {
-        self.context().data().user_data()
+    pub fn data(&self) -> Result<Value, Error> {
+        Ok(self.context()?.data().user_data())
     }
 
     /// @yard
@@ -251,7 +255,7 @@ impl Store {
     /// @return [Integer]
     /// @raise [Error] if fuel consumption is not enabled via {Wasmtime::Engine#new}
     pub fn get_fuel(&self) -> Result<u64, Error> {
-        self.inner_ref().get_fuel().map_err(|e| error!("{}", e))
+        self.inner_ref()?.get_fuel().map_err(|e| error!("{}", e))
     }
 
     /// @yard
@@ -260,7 +264,7 @@ impl Store {
     /// @def set_fuel(fuel)
     /// @raise [Error] if fuel consumption is not enabled via {Wasmtime::Engine#new}
     pub fn set_fuel(&self, fuel: u64) -> Result<(), Error> {
-        unsafe { &mut *self.inner.get() }
+        self.inner_mut()?
             .set_fuel(fuel)
             .map_err(|e| error!("{}", e))?;
 
@@ -277,8 +281,9 @@ impl Store {
     /// @def set_epoch_deadline(ticks_beyond_current)
     /// @param ticks_beyond_current [Integer] The number of ticks before this store reaches the deadline.
     /// @return [nil]
-    pub fn set_epoch_deadline(&self, ticks_beyond_current: u64) {
-        unsafe { &mut *self.inner.get() }.set_epoch_deadline(ticks_beyond_current);
+    pub fn set_epoch_deadline(&self, ticks_beyond_current: u64) -> Result<(), Error> {
+        self.inner_mut()?.set_epoch_deadline(ticks_beyond_current);
+        Ok(())
     }
 
     /// @yard
@@ -286,42 +291,81 @@ impl Store {
     /// Returns whether the linear memory limit has been hit.
     ///
     /// @return [Boolean]
-    pub fn linear_memory_limit_hit(&self) -> bool {
-        self.context().data().store_limits.linear_memory_limit_hit()
+    pub fn linear_memory_limit_hit(&self) -> Result<bool, Error> {
+        Ok(self
+            .context()?
+            .data()
+            .store_limits
+            .linear_memory_limit_hit())
     }
 
     /// @yard
     /// Returns the maximum linear memory consumed.
     ///
     /// @return [Integer]
-    pub fn max_linear_memory_consumed(&self) -> usize {
-        self.context()
+    pub fn max_linear_memory_consumed(&self) -> Result<usize, Error> {
+        Ok(self
+            .context()?
             .data()
             .store_limits
-            .max_linear_memory_consumed()
+            .max_linear_memory_consumed())
     }
 
-    pub fn context(&self) -> StoreContext<'_, StoreData> {
-        unsafe { (*self.inner.get()).as_context() }
+    /// @yard
+    /// Closes the store, immediately freeing its WebAssembly memory.
+    ///
+    /// In Wasmtime the linear memory of every instance lives in the {Store}, so
+    /// this is the way to reclaim that memory deterministically once a job is
+    /// done. Any subsequent use of this store raises an error. The same applies
+    /// to instances, memories, funcs, globals and tables created from it.
+    /// Calling +close+ on an already-closed store is a no-op.
+    ///
+    /// @def close
+    /// @return [Boolean] +true+ if this call closed the store, +false+ if it
+    ///   had already been closed.
+    pub fn close(&self) -> bool {
+        unsafe { (*self.inner.get()).take() }.is_some()
     }
 
-    pub fn context_mut(&self) -> StoreContextMut<'_, StoreData> {
-        unsafe { (*self.inner.get()).as_context_mut() }
+    /// @yard
+    /// @def closed?
+    /// @return [Boolean] Whether this store has been closed via {#close}.
+    pub fn closed(&self) -> bool {
+        unsafe { (*self.inner.get()).is_none() }
+    }
+
+    pub fn context(&self) -> Result<StoreContext<'_, StoreData>, Error> {
+        Ok(self.inner_ref()?.as_context())
+    }
+
+    pub fn context_mut(&self) -> Result<StoreContextMut<'_, StoreData>, Error> {
+        Ok(self.inner_mut()?.as_context_mut())
     }
 
     pub fn retain(&self, value: Value) {
-        self.context_mut().data_mut().retain(value);
+        if let Ok(mut context) = self.context_mut() {
+            context.data_mut().retain(value);
+        }
     }
 
     pub fn take_last_error(&self) -> Option<Error> {
-        self.context_mut().data_mut().take_error()
+        self.context_mut().ok()?.data_mut().take_error()
+    }
+
+    fn inner_ref(&self) -> Result<&StoreImpl<StoreData>, Error> {
+        unsafe { (*self.inner.get()).as_ref() }.ok_or_else(closed_error)
     }
 
-    fn inner_ref(&self) -> &StoreImpl<StoreData> {
-        unsafe { &*self.inner.get() }
+    #[allow(clippy::mut_from_ref)]
+    fn inner_mut(&self) -> Result<&mut StoreImpl<StoreData>, Error> {
+        unsafe { (*self.inner.get()).as_mut() }.ok_or_else(closed_error)
     }
 }
 
+fn closed_error() -> Error {
+    error!("Wasmtime::Store is closed")
+}
+
 /// A wrapper around a Ruby Value that has a store context.
 /// Used in places where both Store or Caller can be used.
 #[derive(Clone, Copy)]
@@ -356,7 +400,7 @@ impl StoreContextValue<'_> {
     pub fn context(&self) -> Result<StoreContext<'_, StoreData>, Error> {
         let ruby = Ruby::get().unwrap();
         match self {
-            Self::Store(store) => Ok(ruby.get_inner_ref(store).context()),
+            Self::Store(store) => ruby.get_inner_ref(store).context(),
             Self::Caller(caller) => ruby.get_inner_ref(caller).context(),
         }
     }
@@ -364,7 +408,7 @@ impl StoreContextValue<'_> {
     pub fn context_mut(&self) -> Result<StoreContextMut<'_, StoreData>, Error> {
         let ruby = Ruby::get().unwrap();
         match self {
-            Self::Store(store) => Ok(ruby.get_inner_ref(store).context_mut()),
+            Self::Store(store) => ruby.get_inner_ref(store).context_mut(),
             Self::Caller(caller) => ruby.get_inner_ref(caller).context_mut(),
         }
     }
@@ -372,11 +416,11 @@ impl StoreContextValue<'_> {
     pub fn set_last_error(&self, error: Error) {
         let ruby = Ruby::get().unwrap();
         match self {
-            Self::Store(store) => ruby
-                .get_inner(*store)
-                .context_mut()
-                .data_mut()
-                .set_error(error),
+            Self::Store(store) => {
+                if let Ok(mut context) = ruby.get_inner(*store).context_mut() {
+                    context.data_mut().set_error(error);
+                }
+            }
             Self::Caller(caller) => {
                 if let Ok(mut context) = ruby.get_inner(*caller).context_mut() {
                     context.data_mut().set_error(error);
@@ -457,6 +501,8 @@ pub fn init(ruby: &Ruby) -> Result<(), Error> {
     class.define_method("get_fuel", method!(Store::get_fuel, 0))?;
     class.define_method("set_fuel", method!(Store::set_fuel, 1))?;
     class.define_method("set_epoch_deadline", method!(Store::set_epoch_deadline, 1))?;
+    class.define_method("close", method!(Store::close, 0))?;
+    class.define_method("closed?", method!(Store::closed, 0))?;
     class.define_method(
         "linear_memory_limit_hit?",
         method!(Store::linear_memory_limit_hit, 0),
diff --git a/ext/src/ruby_api/table.rs b/ext/src/ruby_api/table.rs
index 16bc7059..d653ad79 100644
--- a/ext/src/ruby_api/table.rs
+++ b/ext/src/ruby_api/table.rs
@@ -104,7 +104,7 @@ impl<'a> Table<'a> {
             .clone();
 
         let inner = TableImpl::new(
-            store.context_mut(),
+            store.context_mut()?,
             wasmtime::TableType::new(table_type, min, max),
             ref_,
         )
diff --git a/spec/unit/store_spec.rb b/spec/unit/store_spec.rb
index fc5eba94..01692759 100644
--- a/spec/unit/store_spec.rb
+++ b/spec/unit/store_spec.rb
@@ -162,6 +162,52 @@ module Wasmtime
           GC.enable
         end
       end
+
+      describe "#close" do
+        it "reports closed? before and after" do
+          store = Store.new(engine)
+          expect(store.closed?).to be false
+          store.close
+          expect(store.closed?).to be true
+        end
+
+        it "returns true on first close and false afterwards" do
+          store = Store.new(engine)
+          expect(store.close).to be true
+          expect(store.close).to be false
+        end
+
+        it "closes a store that owns allocated memory" do
+          store = Store.new(engine, limits: {memory_size: 10_000_000})
+          Memory.new(store, min_size: 100)
+          expect(store.close).to be true
+        end
+
+        it "raises when the store is used after close" do
+          store = Store.new(engine)
+          store.close
+          expect { store.data }.to raise_error(Wasmtime::Error, /closed/)
+          expect { store.max_linear_memory_consumed }.to raise_error(Wasmtime::Error, /closed/)
+        end
+
+        it "raises when an instance from a closed store is used" do
+          store = Store.new(engine)
+          mod = Module.new(engine, <<~WAT)
+            (module
+              (func (export "f")))
+          WAT
+          instance = Instance.new(store, mod)
+          store.close
+          expect { instance.invoke("f") }.to raise_error(Wasmtime::Error, /closed/)
+        end
+
+        it "raises when a memory from a closed store is used" do
+          store = Store.new(engine)
+          memory = Memory.new(store, min_size: 1)
+          store.close
+          expect { memory.read(0, 1) }.to raise_error(Wasmtime::Error, /closed/)
+        end
+      end
     end
   end
 end