diff --git a/crates/wasmtime/src/runtime/vm/gc/enabled/copying.rs b/crates/wasmtime/src/runtime/vm/gc/enabled/copying.rs index a602b265f3d3..a648000c12c0 100644 --- a/crates/wasmtime/src/runtime/vm/gc/enabled/copying.rs +++ b/crates/wasmtime/src/runtime/vm/gc/enabled/copying.rs @@ -1185,6 +1185,32 @@ impl GarbageCollection<'_> for CopyingCollection<'_> { } } +impl Drop for CopyingCollection<'_> { + fn drop(&mut self) { + // A collection is logically atomic with respect to the mutator: the + // mutator must never run while the heap is in a half-collected + // (half-flipped, half-forwarded) state. The collection only yields + // between increments to cooperate with the async runtime and + // `Future::poll`; it does not interleave with the mutator. + // + // However, the future driving an incremental collection can be dropped + // before the collection completes (e.g. when a `call_async` future is + // cancelled while a GC is in progress). If we did nothing here, the + // store would be left with a half-collected GC heap and reusing it + // (e.g. starting another `call_async`) would corrupt the GC heap. So, + // to uphold the logical atomicity of collection, we finish any + // in-progress collection here synchronously when dropped. + match self.phase { + CopyingCollectionPhase::ProcessWorklist | CopyingCollectionPhase::SweepExternRefs => { + if let Err(e) = self.collect() { + log::error!("error finishing an interrupted copying collection: {e:?}"); + } + } + CopyingCollectionPhase::ProcessRoots | CopyingCollectionPhase::Done => {} + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/tests/all/gc.rs b/tests/all/gc.rs index b86a73f0b4ae..31468a275c7c 100644 --- a/tests/all/gc.rs +++ b/tests/all/gc.rs @@ -1,8 +1,10 @@ use super::{async_functions::CountPending, gc_store, ref_types_module}; use std::num::NonZeroU32; +use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::Ordering; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering::SeqCst}; +use std::task::{Context, Poll}; use wasmtime::*; struct SetFlagOnDrop(Arc); @@ -3514,3 +3516,123 @@ async fn copying_collector_async_gc_yields() -> Result<()> { Ok(()) } + +// Cancel the wrapped future once a host-triggered collection is observed to be +// `skip` pending-polls underway (so we are past root processing / the semispace +// flip and into the actual copying), simulating an async cancellation that +// drops a `call_async` future while a GC is in progress. +struct CancelDuringGc { + future: F, + in_gc: Arc, + skip: u32, + seen_in_gc: u32, +} + +impl Future for CancelDuringGc { + type Output = Option; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let me = unsafe { self.get_unchecked_mut() }; + let future = unsafe { Pin::new_unchecked(&mut me.future) }; + match future.poll(cx) { + Poll::Ready(val) => Poll::Ready(Some(val)), + Poll::Pending => { + if me.in_gc.load(SeqCst) != 0 { + me.seen_in_gc += 1; + if me.seen_in_gc >= me.skip { + return Poll::Ready(None); + } + } + Poll::Pending + } + } + } +} + +// Test that cancelling a Wasm-execution future mid-collection and then calling +// into Wasm again doesn't trigger or observe any GC heap corruption. +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn issue_13516_copying_collector_cancellation_during_gc() -> Result<()> { + let _ = env_logger::try_init(); + + if crate::no_hog_memory() { + return Ok(()); + } + + let mut config = Config::new(); + config.wasm_gc(true); + config.wasm_function_references(true); + config.collector(Collector::Copying); + + let engine = Engine::new(&config)?; + let mut store = Store::new(&engine, ()); + + // Build a linked list of `N` nodes rooted in a global so a collection has a + // non-trivial worklist to copy (giving a wide window in which to cancel), + // then trigger a host collection. + const N: i32 = 100_000; + let module = Module::new( + &engine, + r#" + (module + (import "" "gc" (func $gc)) + + (type $node (struct (field (ref null $node)))) + + (global $list (mut (ref null $node)) (ref.null $node)) + + (func (export "run") (param $n i32) + (block $done + (loop $loop + (br_if $done (i32.eqz (local.get $n))) + (global.set $list (struct.new $node (global.get $list))) + (local.set $n (i32.sub (local.get $n) (i32.const 1))) + (br $loop) + ) + ) + (call $gc) + ) + ) + "#, + )?; + + // A host function that performs an async collection, flagging while it runs. + let in_gc = Arc::new(AtomicUsize::new(0)); + let mut linker = Linker::new(&engine); + linker.func_wrap_async("", "gc", { + let in_gc = in_gc.clone(); + move |mut caller: Caller<'_, ()>, _: ()| { + let in_gc = in_gc.clone(); + Box::new(async move { + in_gc.fetch_add(1, SeqCst); + let r = caller.gc_async(None).await; + in_gc.store(0, SeqCst); + r + }) + } + })?; + + let instance = linker.instantiate_async(&mut store, &module).await?; + let run = instance.get_typed_func::(&mut store, "run")?; + + // Drive a call that triggers a collection, then cancel it mid-collection. + let cancelled = CancelDuringGc { + future: run.call_async(&mut store, N), + in_gc: in_gc.clone(), + skip: 8, + seen_in_gc: 0, + } + .await + .is_none(); + assert!( + cancelled, + "the collection should have been cancelled mid-flight" + ); + + // Reusing the store after a mid-collection cancellation must not corrupt the + // GC heap. + run.call_async(&mut store, N).await?; + + Ok(()) +}