Fix current instance count on failure during component instantiation

crates/fuzzing/tests/oom/component_instance.rs

@@ -0,0 +1,42 @@
+#![cfg(arc_try_new)]
+
+use wasmtime::component::{Component, Linker};
+use wasmtime::{Config, Engine, PoolingAllocationConfig, Result, Store};
+use wasmtime_fuzzing::oom::OomTest;
+
+#[test]
+fn instantiate() -> Result<()> {
+    let mut config = Config::new();
+    config.concurrency_support(false);
+    let engine = Engine::new(&config)?;
+    let component = Component::new(&engine, "(component)")?;
+    let linker = Linker::<()>::new(&engine);
+    let instance_pre = linker.instantiate_pre(&component)?;
+
+    OomTest::new().test(|| {
+        let mut store = Store::try_new(&engine, ())?;
+        let _instance = instance_pre.instantiate(&mut store)?;
+        Ok(())
+    })
+}
+
+#[test]
+fn instantiate_in_pooling_allocator() -> Result<()> {
+    let mut pool_config = PoolingAllocationConfig::default();
+    pool_config.total_component_instances(1);
+
+    let mut config = Config::new();
+    config.concurrency_support(false);
+    config.allocation_strategy(pool_config);
+
+    let engine = Engine::new(&config)?;
+    let component = Component::new(&engine, "(component)")?;
+    let linker = Linker::<()>::new(&engine);
+    let instance_pre = linker.instantiate_pre(&component)?;
+
+    OomTest::new().test(|| {
+        let mut store = Store::try_new(&engine, ())?;
+        let _instance = instance_pre.instantiate(&mut store)?;
+        Ok(())
+    })
+}

crates/fuzzing/tests/oom/main.rs

@@ -6,6 +6,7 @@ mod boxed;
 mod btree_map;
 mod caller;
 mod component_func;
+mod component_instance;
 mod component_linker;
 mod component_resource;
 mod config;

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

@@ -1187,23 +1187,46 @@ impl<T: 'static> InstancePre<T> {
         mut store: impl AsContextMut<Data = T>,
         asyncness: Asyncness,
     ) -> Result<Instance> {
-        let mut store = store.as_context_mut();
+        let store = store.as_context_mut();
         store.0.set_async_required(self.asyncness);
         store
             .engine()
             .allocator()
             .increment_component_instance_count()?;
-        let mut instantiator = Instantiator::new(&self.component, store.0, &self.imports)?;
-        instantiator.run(&mut store, asyncness).await.map_err(|e| {
-            store
-                .engine()
-                .allocator()
-                .decrement_component_instance_count();
-            e
-        })?;
-
-        let instance = Instance::from_wasmtime(store.0, instantiator.id);
-        store.0.push_component_instance(instance);
-        Ok(instance)
+
+        // Helper structure to pair the above increment with a decrement should
+        // anything fail below.
+        let mut decrement = DecrementComponentInstanceCountOnDrop {
+            store,
+            enabled: true,
+        };
+
+        let mut instantiator =
+            Instantiator::new(&self.component, decrement.store.0, &self.imports)?;
+        instantiator.run(&mut decrement.store, asyncness).await?;
+
+        let instance = Instance::from_wasmtime(decrement.store.0, instantiator.id);
+        decrement.store.0.push_component_instance(instance);
+
+        // Everything has passed, don't decrement the instance count and let the
+        // destructor for the `Store` handle that at this point.
+        decrement.enabled = false;
+        return Ok(instance);
+
+        struct DecrementComponentInstanceCountOnDrop<'a, T: 'static> {
+            store: StoreContextMut<'a, T>,
+            enabled: bool,
+        }
+
+        impl<T> Drop for DecrementComponentInstanceCountOnDrop<'_, T> {
+            fn drop(&mut self) {
+                if self.enabled {
+                    self.store
+                        .engine()
+                        .allocator()
+                        .decrement_component_instance_count();
+                }
+            }
+        }
     }
 }