fix memory leak when panicking with detached ComponentInstance

In `ComponentInstance::from_vmctx` and
`StoreContextMut::with_detached_instance[_async]` we were leaking memory due to
having taken `InstanceData` out of the `Store` and making it unreachable except
via a raw pointer, meaning there was nothing responsible for dropping it on
panic.  This commit stashes the data elsewhere in the store to make sure it gets
dropped.

Signed-off-by: Joel Dice <joel.dice@fermyon.com>

Author: dicej

crates/wasmtime/src/runtime/component/concurrent.rs

@@ -715,21 +715,14 @@ impl<T> StoreContextMut<'_, T> {
             }
             _ => unreachable!(),
         }));
-        let data = self.0[instance.0].take().unwrap();
-        let (result, data) = {
-            // SAFETY: We've taken the instance out of the store, so now we own
-            // it and can take an exclusive reference to it.
-            let instance = unsafe { &mut *data.instance_ptr() };
-            assert!(instance.data.is_none());
-            instance.data = Some(data);
-            instance.set_store(None);
-            let result = fun(self.as_context_mut(), instance);
-            instance.set_store(Some(VMStoreRawPtr(self.traitobj())));
-            (result, instance.data.take())
-        };
-        if self.0[instance.0].is_none() {
-            self.0[instance.0] = data;
-        }
+        let ptr = self.0.hide_instance(*instance);
+        // SAFETY: We've taken the instance out of the store, so now we own
+        // it and can take an exclusive reference to it.
+        let reference = unsafe { &mut *ptr };
+        reference.set_store(None);
+        let result = fun(self.as_context_mut(), reference);
+        reference.set_store(Some(VMStoreRawPtr(self.traitobj())));
+        self.0.unhide_instance(*instance);
         result
     }
 
@@ -746,21 +739,14 @@ impl<T> StoreContextMut<'_, T> {
             }
             _ => unreachable!(),
         }));
-        let data = self.0[instance.0].take().unwrap();
-        let (result, data) = {
-            // SAFETY: We've taken the instance out of the store, so now we own
-            // it and can take an exclusive reference to it.
-            let instance = unsafe { &mut *data.instance_ptr() };
-            assert!(instance.data.is_none());
-            instance.data = Some(data);
-            instance.set_store(None);
-            let result = fun(self.as_context_mut(), instance).await;
-            instance.set_store(Some(VMStoreRawPtr(self.traitobj())));
-            (result, instance.data.take())
-        };
-        if self.0[instance.0].is_none() {
-            self.0[instance.0] = data;
-        }
+        let ptr = self.0.hide_instance(*instance);
+        // SAFETY: We've taken the instance out of the store, so now we own
+        // it and can take an exclusive reference to it.
+        let reference = unsafe { &mut *ptr };
+        reference.set_store(None);
+        let result = fun(self.as_context_mut(), reference).await;
+        reference.set_store(Some(VMStoreRawPtr(self.traitobj())));
+        self.0.unhide_instance(*instance);
         result
     }
 
@@ -783,13 +769,13 @@ impl<T> StoreContextMut<'_, T> {
             _ => unreachable!(),
         }));
         if let Some(handle) = instance.instance {
-            self.0[handle.0] = Some(instance.data.take().unwrap());
+            self.0.unhide_instance(handle);
         }
         instance.set_store(Some(VMStoreRawPtr(self.traitobj())));
         let result = fun(self.as_context_mut(), instance.instance);
         instance.set_store(None);
         if let Some(handle) = instance.instance {
-            instance.data = Some(self.0[handle.0].take().unwrap());
+            self.0.hide_instance(handle);
         }
         result
     }

crates/wasmtime/src/runtime/store.rs

@@ -410,6 +410,9 @@ pub struct StoreOpaque {
     #[cfg(feature = "component-model-async")]
     concurrent_async_state: concurrent::AsyncState,
 
+    #[cfg(feature = "component-model")]
+    hidden_instances: Vec<Option<Box<crate::component::InstanceData>>>,
+
     /// State related to the executor of wasm code.
     ///
     /// For example if Pulley is enabled and configured then this will store a
@@ -587,6 +590,8 @@ impl<T: 'static> Store<T> {
             component_calls: Default::default(),
             #[cfg(feature = "component-model")]
             host_resource_data: Default::default(),
+            #[cfg(feature = "component-model")]
+            hidden_instances: Default::default(),
             #[cfg(feature = "component-model-async")]
             concurrent_async_state: Default::default(),
             #[cfg(has_host_compiler_backend)]
@@ -1183,6 +1188,29 @@ impl StoreOpaque {
         self.store_data.id()
     }
 
+    #[cfg(feature = "component-model")]
+    pub(crate) fn hide_instance(
+        &mut self,
+        handle: crate::component::Instance,
+    ) -> *mut crate::vm::component::ComponentInstance {
+        if self.hidden_instances.len() <= handle.0.index() {
+            self.hidden_instances
+                .resize_with(handle.0.index() + 1, || None);
+        }
+        let data = self[handle.0].take().unwrap();
+        let ptr = data.instance_ptr();
+        self.hidden_instances[handle.0.index()] = Some(data);
+        ptr
+    }
+
+    #[cfg(feature = "component-model")]
+    pub(crate) fn unhide_instance(&mut self, handle: crate::component::Instance) {
+        if let Some(data) = self.hidden_instances[handle.0.index()].take() {
+            assert!(self[handle.0].is_none());
+            self[handle.0] = Some(data);
+        }
+    }
+
     pub fn bump_resource_counts(&mut self, module: &Module) -> Result<()> {
         fn bump(slot: &mut usize, max: usize, amt: usize, desc: &str) -> Result<()> {
             let new = slot.saturating_add(amt);

crates/wasmtime/src/runtime/store/data.rs

@@ -296,7 +296,7 @@ impl<T> Stored<T> {
         self.store_id.assert_belongs_to(store)
     }
 
-    fn index(&self) -> usize {
+    pub(crate) fn index(&self) -> usize {
         self.index
     }
 }

crates/wasmtime/src/runtime/vm/component.rs

@@ -6,7 +6,7 @@
 //! Eventually it's intended that module-to-module calls, which would be
 //! cranelift-compiled adapters, will use this `VMComponentContext` as well.
 
-use crate::component::{Instance, InstanceData, ResourceType};
+use crate::component::{Instance, ResourceType};
 use crate::prelude::*;
 use crate::runtime::vm::{
     SendSyncPtr, VMArrayCallFunction, VMContext, VMFuncRef, VMGlobalDefinition, VMMemoryDefinition,
@@ -78,8 +78,6 @@ pub struct ComponentInstance {
 
     pub(crate) instance: Option<Instance>,
 
-    pub(crate) data: Option<Box<InstanceData>>,
-
     /// A zero-sized field which represents the end of the struct for the actual
     /// `VMComponentContext` to be allocated behind.
     vmctx: VMComponentContext,
@@ -189,16 +187,13 @@ impl ComponentInstance {
         let reference = ptr.as_mut();
         let store = &mut *reference.store();
         if let Some(instance) = reference.instance {
-            assert!(reference.data.is_none());
-            reference.data = Some(store[instance.0].take().unwrap());
+            store.hide_instance(instance);
         }
         reference.set_store(None);
         let result = f(store, reference);
         reference.set_store(Some(VMStoreRawPtr(store.traitobj())));
         if let Some(instance) = reference.instance {
-            if store[instance.0].is_none() {
-                store[instance.0] = reference.data.take();
-            }
+            store.unhide_instance(instance);
         }
         result
     }
@@ -265,7 +260,6 @@ impl ComponentInstance {
                 vmctx: VMComponentContext {
                     _marker: marker::PhantomPinned,
                 },
-                data: None,
                 instance: None,
                 #[cfg(feature = "component-model-async")]
                 concurrent_state,