Repository: bytecodealliance/wasmtime
Pull request #6509 (2023-06-02 16:03:10)

Fix a soundness issue with the component model and async
This commit addresses https://github.com/bytecodealliance/wasmtime/issues/6493 by fixing a soundness issue with the async
implementation of the component model. This issue has been presence
since the inception of the addition of async support to the component
model and doesn't represent a recent regression. The underlying problem
is that one of the base assumptions of the trap handling code is that
there's only one single activation in TLS that needs to be pushed/popped
when a stack is switched (e.g. a fiber is switched to or from). In the
case of the component model there might be two activations: one for an
invocation of a component function and then a second for an invocation
of a `realloc` function to return results back to wasm (e.g. in the case
an imported function returns a list).

This problem is fixed by changing how TLS is managed in the presence of
fibers. Previously when a fiber was suspended it would pop a single
activation from the top of the stack and save that to get pushed when
the fiber was resumed. This has the benefit of maintaining an entire
linked list of activations for the current thread but has the problem
above where it doesn't handle a fiber with multiple activations on it.
Instead now TLS management is done when a fiber is resumed instead of
suspended. Instead of pushing/popping a single activation the entire
linked list of activations is tracked for a particular fiber and stored
within the fiber itself. In this manner resuming a fiber will push
all activations onto the current thread and suspending a fiber will pop
all activations for the fiber (and store them as a new linked list in
the fiber's state itself).

This end result is that all activations on a fiber should now be managed
correctly, regardless of how many there are. The main downside of this
commit is that fiber suspension and resumption is more complicated, but
the hope there is that fiber suspension typically corresponds with I/O
not being ready or similar so the order of magnitude of TLS operations
isn't too significant compared to the I/O overhead.

alexcrichton: I'll note that I had a number of false starts in attempting to solve this issue. First I thought that this only needed a push/pop around calls to `cabi_realloc` in the component model, but that wasn't correct because that only works for import-based calls to `realloc`, not export-based calls. I then wrote a test to expose the flaw with that strategy.

Next I thought that it would be possible to not actually maintain the entire linked list of activations and instead have a fiber have its own linked list which replaces the current thread's linked list when it resumes and swaps it back out when it suspends. This however does not work because a complete list of activations is required to correctly perform GC, so I wrote a test to expose that bug.

Finally I settled on this solution, which is at least sufficient for all the above scenarios. I'm not 100% sure it's the end-all-be-all, so extra care in review would be appreciated!

Patch:
diff --git a/Cargo.lock b/Cargo.lock
index b1b0abca1..bd5833de9 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1303,9 +1303,9 @@ dependencies = [
 
 [[package]]
 name = "file-per-thread-logger"
-version = "0.1.6"
+version = "0.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "84f2e425d9790201ba4af4630191feac6dcc98765b118d4d18e91d23c2353866"
+checksum = "8a3cc21c33af89af0930c8cae4ade5e6fdc17b5d2c97b3d2e2edb67a1cf683f3"
 dependencies = [
  "env_logger 0.10.0",
  "log",
diff --git a/Cargo.toml b/Cargo.toml
index 0aacb7efa..75d415493 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -238,6 +238,7 @@ cfg-if = "1.0"
 tempfile = "3.1.0"
 filecheck = "0.5.0"
 libc = "0.2.60"
+file-per-thread-logger = "0.2.0"
 
 [features]
 default = [
diff --git a/RELEASES.md b/RELEASES.md
index 64e1b2d11..7f1fc7c6f 100644
--- a/RELEASES.md
+++ b/RELEASES.md
@@ -12,6 +12,10 @@ Unreleased.
 
 ### Changed
 
+* An `epoch_deadline_callback` now returns an `UpdateDeadline` enum to allow
+  optionally yielding to the async executor after the callback runs.
+  [#6464](https://github.com/bytecodealliance/wasmtime/pull/6464)
+
 * The "raw" representation of `funcref` and `externref` in the embedding API has
   been updated from a `usize` to a `*mut u8` to be compatible with Rust's
   proposed strict provenance rules. This change is additionally reflected into
diff --git a/crates/c-api/src/store.rs b/crates/c-api/src/store.rs
index 5aedaf371..d56977619 100644
--- a/crates/c-api/src/store.rs
+++ b/crates/c-api/src/store.rs
@@ -4,7 +4,7 @@ use std::ffi::c_void;
 use std::sync::Arc;
 use wasmtime::{
     AsContext, AsContextMut, Store, StoreContext, StoreContextMut, StoreLimits, StoreLimitsBuilder,
-    Val,
+    UpdateDeadline, Val,
 };
 
 /// This representation of a `Store` is used to implement the `wasm.h` API.
@@ -143,7 +143,7 @@ pub extern "C" fn wasmtime_store_epoch_deadline_callback(
             Some(err) => Err(wasmtime::Error::from(<wasmtime_error_t as Into<
                 anyhow::Error,
             >>::into(*err))),
-            None => Ok(delta),
+            None => Ok(UpdateDeadline::Continue(delta)),
         }
     });
 }
diff --git a/crates/cache/Cargo.toml b/crates/cache/Cargo.toml
index b7fb05988..6ff7e9884 100644
--- a/crates/cache/Cargo.toml
+++ b/crates/cache/Cargo.toml
@@ -13,7 +13,7 @@ anyhow = { workspace = true }
 base64 = "0.21.0"
 bincode = "1.1.4"
 directories-next = "2.0"
-file-per-thread-logger = "0.1.1"
+file-per-thread-logger = { workspace = true }
 log = { workspace = true }
 serde = { version = "1.0.94", features = ["derive"] }
 sha2 = "0.10.2"
diff --git a/crates/cli-flags/Cargo.toml b/crates/cli-flags/Cargo.toml
index 45bf2b748..38a08f02b 100644
--- a/crates/cli-flags/Cargo.toml
+++ b/crates/cli-flags/Cargo.toml
@@ -11,7 +11,7 @@ edition.workspace = true
 [dependencies]
 anyhow = { workspace = true }
 clap = { workspace = true }
-file-per-thread-logger = "0.1.1"
+file-per-thread-logger = { workspace = true }
 pretty_env_logger = "0.4.0"
 rayon = "1.5.0"
 wasmtime = { workspace = true }
diff --git a/crates/fiber/src/lib.rs b/crates/fiber/src/lib.rs
index 3e159ee70..79d7bdaa6 100644
--- a/crates/fiber/src/lib.rs
+++ b/crates/fiber/src/lib.rs
@@ -2,6 +2,7 @@ use std::any::Any;
 use std::cell::Cell;
 use std::io;
 use std::marker::PhantomData;
+use std::ops::Range;
 use std::panic::{self, AssertUnwindSafe};
 
 cfg_if::cfg_if! {
@@ -26,24 +27,35 @@ impl FiberStack {
         Ok(Self(imp::FiberStack::new(size)?))
     }
 
-    /// Creates a new fiber stack with the given pointer to the top of the stack.
+    /// Creates a new fiber stack with the given pointer to the bottom of the
+    /// stack plus the byte length of the stack.
+    ///
+    /// The `bottom` pointer should be addressable for `len` bytes. The page
+    /// beneath `bottom` should be unmapped as a guard page.
     ///
     /// # Safety
     ///
-    /// This is unsafe because there is no validation of the given stack pointer.
+    /// This is unsafe because there is no validation of the given pointer.
     ///
     /// The caller must properly allocate the stack space with a guard page and
     /// make the pages accessible for correct behavior.
-    pub unsafe fn from_top_ptr(top: *mut u8) -> io::Result<Self> {
-        Ok(Self(imp::FiberStack::from_top_ptr(top)?))
+    pub unsafe fn from_raw_parts(bottom: *mut u8, len: usize) -> io::Result<Self> {
+        Ok(Self(imp::FiberStack::from_raw_parts(bottom, len)?))
     }
 
     /// Gets the top of the stack.
     ///
-    /// Returns `None` if the platform does not support getting the top of the stack.
+    /// Returns `None` if the platform does not support getting the top of the
+    /// stack.
     pub fn top(&self) -> Option<*mut u8> {
         self.0.top()
     }
+
+    /// Returns the range of where this stack resides in memory if the platform
+    /// supports it.
+    pub fn range(&self) -> Option<Range<usize>> {
+        self.0.range()
+    }
 }
 
 pub struct Fiber<'a, Resume, Yield, Return> {
diff --git a/crates/fiber/src/unix.rs b/crates/fiber/src/unix.rs
index e9c4ce64f..ef0ddea6b 100644
--- a/crates/fiber/src/unix.rs
+++ b/crates/fiber/src/unix.rs
@@ -34,6 +34,7 @@
 use crate::RunResult;
 use std::cell::Cell;
 use std::io;
+use std::ops::Range;
 use std::ptr;
 
 #[derive(Debug)]
@@ -41,8 +42,10 @@ pub struct FiberStack {
     // The top of the stack; for stacks allocated by the fiber implementation itself,
     // the base address of the allocation will be `top.sub(len.unwrap())`
     top: *mut u8,
-    // The length of the stack; `None` when the stack was not created by this implementation.
-    len: Option<usize>,
+    // The length of the stack
+    len: usize,
+    // whether or not this stack was mmap'd
+    mmap: bool,
 }
 
 impl FiberStack {
@@ -74,25 +77,35 @@ impl FiberStack {
 
             Ok(Self {
                 top: mmap.cast::<u8>().add(mmap_len),
-                len: Some(mmap_len),
+                len: mmap_len,
+                mmap: true,
             })
         }
     }
 
-    pub unsafe fn from_top_ptr(top: *mut u8) -> io::Result<Self> {
-        Ok(Self { top, len: None })
+    pub unsafe fn from_raw_parts(base: *mut u8, len: usize) -> io::Result<Self> {
+        Ok(Self {
+            top: base.add(len),
+            len,
+            mmap: false,
+        })
     }
 
     pub fn top(&self) -> Option<*mut u8> {
         Some(self.top)
     }
+
+    pub fn range(&self) -> Option<Range<usize>> {
+        let base = unsafe { self.top.sub(self.len) as usize };
+        Some(base..base + self.len)
+    }
 }
 
 impl Drop for FiberStack {
     fn drop(&mut self) {
         unsafe {
-            if let Some(len) = self.len {
-                let ret = rustix::mm::munmap(self.top.sub(len) as _, len);
+            if self.mmap {
+                let ret = rustix::mm::munmap(self.top.sub(self.len) as _, self.len);
                 debug_assert!(ret.is_ok());
             }
         }
diff --git a/crates/fiber/src/windows.rs b/crates/fiber/src/windows.rs
index 5b9ada3cc..04452bd87 100644
--- a/crates/fiber/src/windows.rs
+++ b/crates/fiber/src/windows.rs
@@ -2,6 +2,7 @@ use crate::RunResult;
 use std::cell::Cell;
 use std::ffi::c_void;
 use std::io;
+use std::ops::Range;
 use std::ptr;
 use windows_sys::Win32::Foundation::*;
 use windows_sys::Win32::System::Threading::*;
@@ -14,13 +15,17 @@ impl FiberStack {
         Ok(Self(size))
     }
 
-    pub unsafe fn from_top_ptr(_top: *mut u8) -> io::Result<Self> {
+    pub unsafe fn from_raw_parts(_base: *mut u8, _len: usize) -> io::Result<Self> {
         Err(io::Error::from_raw_os_error(ERROR_NOT_SUPPORTED as i32))
     }
 
     pub fn top(&self) -> Option<*mut u8> {
         None
     }
+
+    pub fn range(&self) -> Option<Range<usize>> {
+        None
+    }
 }
 
 pub struct Fiber {
diff --git a/crates/runtime/src/instance/allocator/pooling.rs b/crates/runtime/src/instance/allocator/pooling.rs
index 5d3d57bbe..fd99c0f3b 100644
--- a/crates/runtime/src/instance/allocator/pooling.rs
+++ b/crates/runtime/src/instance/allocator/pooling.rs
@@ -462,7 +462,7 @@ impl StackPool {
             commit_stack_pages(bottom_of_stack, size_without_guard)?;
 
             let stack =
-                wasmtime_fiber::FiberStack::from_top_ptr(bottom_of_stack.add(size_without_guard))?;
+                wasmtime_fiber::FiberStack::from_raw_parts(bottom_of_stack, size_without_guard)?;
             Ok(stack)
         }
     }
diff --git a/crates/runtime/src/lib.rs b/crates/runtime/src/lib.rs
index 8ae55575f..e23851799 100644
--- a/crates/runtime/src/lib.rs
+++ b/crates/runtime/src/lib.rs
@@ -69,10 +69,7 @@ pub use crate::mmap::Mmap;
 pub use crate::mmap_vec::MmapVec;
 pub use crate::store_box::*;
 pub use crate::table::{Table, TableElement};
-pub use crate::traphandlers::{
-    catch_traps, init_traps, raise_lib_trap, raise_user_trap, resume_panic, tls_eager_initialize,
-    Backtrace, Frame, SignalHandler, TlsRestore, Trap, TrapReason,
-};
+pub use crate::traphandlers::*;
 pub use crate::vmcontext::{
     VMArrayCallFunction, VMArrayCallHostFuncContext, VMContext, VMFuncRef, VMFunctionBody,
     VMFunctionImport, VMGlobalDefinition, VMGlobalImport, VMInvokeArgument, VMMemoryDefinition,
diff --git a/crates/runtime/src/traphandlers.rs b/crates/runtime/src/traphandlers.rs
index 8dad271cf..25d852671 100644
--- a/crates/runtime/src/traphandlers.rs
+++ b/crates/runtime/src/traphandlers.rs
@@ -12,7 +12,7 @@ use std::ptr;
 use std::sync::Once;
 
 pub use self::backtrace::{Backtrace, Frame};
-pub use self::tls::{tls_eager_initialize, TlsRestore};
+pub use self::tls::{tls_eager_initialize, AsyncWasmCallState, PreviousAsyncWasmCallState};
 
 cfg_if::cfg_if! {
     if #[cfg(miri)] {
@@ -297,7 +297,7 @@ mod call_thread_state {
 
         pub(crate) limits: *const VMRuntimeLimits,
 
-        prev: Cell<tls::Ptr>,
+        pub(super) prev: Cell<tls::Ptr>,
 
         // The values of `VMRuntimeLimits::last_wasm_{exit_{pc,fp},entry_sp}`
         // for the *previous* `CallThreadState` for this same store/limits. Our
@@ -517,6 +517,8 @@ impl<T: Copy> Drop for ResetCell<'_, T> {
 // the caller to the trap site.
 mod tls {
     use super::CallThreadState;
+    use std::mem;
+    use std::ops::Range;
 
     pub use raw::Ptr;
 
@@ -588,49 +590,132 @@ mod tls {
 
     pub use raw::initialize as tls_eager_initialize;
 
-    /// Opaque state used to help control TLS state across stack switches for
-    /// async support.
-    pub struct TlsRestore {
+    /// Opaque state used to persist the state of the `CallThreadState`
+    /// activations associated with a fiber stack that's used as part of an
+    /// async wasm call.
+    pub struct AsyncWasmCallState {
+        // The head of a linked list of activations that are currently present
+        // on an async call's fiber stack. This pointer points to the oldest
+        // activation frame where the `prev` links internally link to younger
+        // activation frames.
+        //
+        // When pushed onto a thread this linked list is traversed to get pushed
+        // onto the current thread at the time.
         state: raw::Ptr,
     }
 
-    impl TlsRestore {
-        /// Takes the TLS state that is currently configured and returns a
-        /// token that is used to replace it later.
+    impl AsyncWasmCallState {
+        /// Creates new state that initially starts as null.
+        pub fn new() -> AsyncWasmCallState {
+            AsyncWasmCallState {
+                state: std::ptr::null_mut(),
+            }
+        }
+
+        /// Pushes the saved state of this wasm's call onto the current thread's
+        /// state.
+        ///
+        /// This will iterate over the linked list of states stored within
+        /// `self` and push them sequentially onto the current thread's
+        /// activation list.
+        ///
+        /// The returned `PreviousAsyncWasmCallState` captures the state of this
+        /// thread just before this operation, and it must have its `restore`
+        /// method called to restore the state when the async wasm is suspended
+        /// from.
         ///
-        /// This is not a safe operation since it's intended to only be used
-        /// with stack switching found with fibers and async wasmtime.
-        pub unsafe fn take() -> TlsRestore {
-            // Our tls pointer must be set at this time, and it must not be
-            // null. We need to restore the previous pointer since we're
-            // removing ourselves from the call-stack, and in the process we
-            // null out our own previous field for safety in case it's
-            // accidentally used later.
-            let state = raw::get();
-            if let Some(state) = state.as_ref() {
-                state.pop();
-            } else {
-                // Null case: we aren't in a wasm context, so theres no tls to
-                // save for restoration.
+        /// # Unsafety
+        ///
+        /// Must be carefully coordinated with
+        /// `PreviousAsyncWasmCallState::restore` and fiber switches to ensure
+        /// that this doesn't push stale data and the data is popped
+        /// appropriately.
+        pub unsafe fn push(self) -> PreviousAsyncWasmCallState {
+            // Our `state` pointer is a linked list of oldest-to-youngest so by
+            // pushing in order of the list we restore the youngest-to-oldest
+            // list as stored in the state of this current thread.
+            let ret = PreviousAsyncWasmCallState { state: raw::get() };
+            let mut ptr = self.state;
+            while let Some(state) = ptr.as_ref() {
+                ptr = state.prev.replace(std::ptr::null_mut());
+                state.push();
             }
+            ret
+        }
 
-            TlsRestore { state }
+        /// Performs a runtime check that this state is indeed null.
+        pub fn assert_null(&self) {
+            assert!(self.state.is_null());
         }
 
-        /// Restores a previous tls state back into this thread's TLS.
+        /// Asserts that the current CallThreadState pointer, if present, is not
+        /// in the `range` specified.
         ///
-        /// This is unsafe because it's intended to only be used within the
-        /// context of stack switching within wasmtime.
-        pub unsafe fn replace(self) {
-            if let Some(state) = self.state.as_ref() {
-                state.push();
-            } else {
-                // Null case: we aren't in a wasm context, so theres no tls
-                // to restore.
+        /// This is used when exiting a future in Wasmtime to assert that the
+        /// current CallThreadState pointer does not point within the stack
+        /// we're leaving (e.g.  allocated for a fiber).
+        pub fn assert_current_state_not_in_range(range: Range<usize>) {
+            let p = raw::get() as usize;
+            assert!(p < range.start || range.end < p);
+        }
+    }
+
+    /// Opaque state used to help control TLS state across stack switches for
+    /// async support.
+    pub struct PreviousAsyncWasmCallState {
+        // The head of a linked list, similar to the TLS state. Note though that
+        // this list is stored in reverse order to assist with `push` and `pop`
+        // below.
+        //
+        // After a `push` call this stores the previous head for the current
+        // thread so we know when to stop popping during a `pop`.
+        state: raw::Ptr,
+    }
+
+    impl PreviousAsyncWasmCallState {
+        /// Pops a fiber's linked list of activations and stores them in
+        /// `AsyncWasmCallState`.
+        ///
+        /// This will pop the top activation of this current thread continuously
+        /// until it reaches whatever the current activation was when `push` was
+        /// originally called.
+        ///
+        /// # Unsafety
+        ///
+        /// Must be paired with a `push` and only performed at a time when a
+        /// fiber is being suspended.
+        pub unsafe fn restore(self) -> AsyncWasmCallState {
+            let thread_head = self.state;
+            mem::forget(self);
+            let mut ret = AsyncWasmCallState::new();
+            loop {
+                // If the current TLS state is as we originally found it, then
+                // this loop is finished.
+                let ptr = raw::get();
+                if ptr == thread_head {
+                    break ret;
+                }
+
+                // Pop this activation from the current thread's TLS state, and
+                // then afterwards push it onto our own linked list within this
+                // `AsyncWasmCallState`. Note that the linked list in `AsyncWasmCallState` is stored
+                // in reverse order so a subsequent `push` later on pushes
+                // everything in the right order.
+                (*ptr).pop();
+                if let Some(state) = ret.state.as_ref() {
+                    (*ptr).prev.set(state);
+                }
+                ret.state = ptr;
             }
         }
     }
 
+    impl Drop for PreviousAsyncWasmCallState {
+        fn drop(&mut self) {
+            panic!("must be consumed with `restore`");
+        }
+    }
+
     /// Configures thread local state such that for the duration of the
     /// execution of `closure` any call to `with` will yield `state`, unless
     /// this is recursively called again.
diff --git a/crates/wasmtime/src/lib.rs b/crates/wasmtime/src/lib.rs
index 2dfac02c2..37e581852 100644
--- a/crates/wasmtime/src/lib.rs
+++ b/crates/wasmtime/src/lib.rs
@@ -424,7 +424,9 @@ pub use crate::profiling::GuestProfiler;
 pub use crate::r#ref::ExternRef;
 #[cfg(feature = "async")]
 pub use crate::store::CallHookHandler;
-pub use crate::store::{AsContext, AsContextMut, CallHook, Store, StoreContext, StoreContextMut};
+pub use crate::store::{
+    AsContext, AsContextMut, CallHook, Store, StoreContext, StoreContextMut, UpdateDeadline,
+};
 pub use crate::trap::*;
 pub use crate::types::*;
 pub use crate::values::*;
diff --git a/crates/wasmtime/src/store.rs b/crates/wasmtime/src/store.rs
index 14b3a3d63..930a73ee9 100644
--- a/crates/wasmtime/src/store.rs
+++ b/crates/wasmtime/src/store.rs
@@ -204,7 +204,8 @@ pub struct StoreInner<T> {
 
     limiter: Option<ResourceLimiterInner<T>>,
     call_hook: Option<CallHookInner<T>>,
-    epoch_deadline_behavior: EpochDeadline<T>,
+    epoch_deadline_behavior:
+        Option<Box<dyn FnMut(StoreContextMut<T>) -> Result<UpdateDeadline> + Send + Sync>>,
     // for comments about `ManuallyDrop`, see `Store::into_data`
     data: ManuallyDrop<T>,
 }
@@ -230,6 +231,18 @@ enum CallHookInner<T> {
     Async(Box<dyn CallHookHandler<T> + Send + Sync>),
 }
 
+/// What to do after returning from a callback when the engine epoch reaches
+/// the deadline for a Store during execution of a function using that store.
+pub enum UpdateDeadline {
+    /// Extend the deadline by the specified number of ticks.
+    Continue(u64),
+    /// Extend the deadline by the specified number of ticks after yielding to
+    /// the async executor loop. This can only be used with an async [`Store`]
+    /// configured via [`Config::async_support`](crate::Config::async_support).
+    #[cfg(feature = "async")]
+    Yield(u64),
+}
+
 // Forward methods on `StoreOpaque` to also being on `StoreInner<T>`
 impl<T> Deref for StoreInner<T> {
     type Target = StoreOpaque;
@@ -422,21 +435,6 @@ enum OutOfGas {
     },
 }
 
-/// What to do when the engine epoch reaches the deadline for a Store
-/// during execution of a function using that store.
-#[derive(Default)]
-enum EpochDeadline<T> {
-    /// Return early with a trap.
-    #[default]
-    Trap,
-    /// Call a custom deadline handler.
-    Callback(Box<dyn FnMut(StoreContextMut<T>) -> Result<u64> + Send + Sync>),
-    /// Extend the deadline by the specified number of ticks after
-    /// yielding to the async executor loop.
-    #[cfg(feature = "async")]
-    YieldAndExtendDeadline { delta: u64 },
-}
-
 impl<T> Store<T> {
     /// Creates a new [`Store`] to be associated with the given [`Engine`] and
     /// `data` provided.
@@ -480,7 +478,7 @@ impl<T> Store<T> {
             },
             limiter: None,
             call_hook: None,
-            epoch_deadline_behavior: EpochDeadline::Trap,
+            epoch_deadline_behavior: None,
             data: ManuallyDrop::new(data),
         });
 
@@ -928,10 +926,17 @@ impl<T> Store<T> {
     /// store and the epoch deadline is reached before completion, the
     /// provided callback function is invoked.
     ///
-    /// This function should return a positive `delta`, which is used to
-    /// update the new epoch, setting it to the current epoch plus
-    /// `delta` ticks. Alternatively, the callback may return an error,
-    /// which will terminate execution.
+    /// This callback should either return an [`UpdateDeadline`], or
+    /// return an error, which will terminate execution with a trap.
+    ///
+    /// The [`UpdateDeadline`] is a positive number of ticks to
+    /// add to the epoch deadline, as well as indicating what
+    /// to do after the callback returns. If the [`Store`] is
+    /// configured with async support, then the callback may return
+    /// [`UpdateDeadline::Yield`] to yield to the async executor before
+    /// updating the epoch deadline. Alternatively, the callback may
+    /// return [`UpdateDeadline::Continue`] to update the epoch deadline
+    /// immediately.
     ///
     /// This setting is intended to allow for coarse-grained
     /// interruption, but not a deterministic deadline of a fixed,
@@ -943,7 +948,7 @@ impl<T> Store<T> {
     /// for an introduction to epoch-based interruption.
     pub fn epoch_deadline_callback(
         &mut self,
-        callback: impl FnMut(StoreContextMut<T>) -> Result<u64> + Send + Sync + 'static,
+        callback: impl FnMut(StoreContextMut<T>) -> Result<UpdateDeadline> + Send + Sync + 'static,
     ) {
         self.inner.epoch_deadline_callback(Box::new(callback));
     }
@@ -1593,6 +1598,7 @@ impl<T> StoreContextMut<'_, T> {
                 fiber,
                 current_poll_cx,
                 engine,
+                state: Some(wasmtime_runtime::AsyncWasmCallState::new()),
             }
         };
         future.await?;
@@ -1603,6 +1609,8 @@ impl<T> StoreContextMut<'_, T> {
             fiber: wasmtime_fiber::Fiber<'a, Result<()>, (), Result<()>>,
             current_poll_cx: *mut *mut Context<'static>,
             engine: Engine,
+            // See comments in `FiberFuture::resume` for this
+            state: Option<wasmtime_runtime::AsyncWasmCallState>,
         }
 
         // This is surely the most dangerous `unsafe impl Send` in the entire
@@ -1668,6 +1676,50 @@ impl<T> StoreContextMut<'_, T> {
         // correct. That's what `unsafe` in Rust is all about, though, right?
         unsafe impl Send for FiberFuture<'_> {}
 
+        impl FiberFuture<'_> {
+            /// This is a helper function to call `resume` on the underlying
+            /// fiber while correctly managing Wasmtime's thread-local data.
+            ///
+            /// Wasmtime's implementation of traps leverages thread-local data
+            /// to get access to metadata during a signal. This thread-local
+            /// data is a linked list of "activations" where the nodes of the
+            /// linked list are stored on the stack. It would be invalid as a
+            /// result to suspend a computation with the head of the linked list
+            /// on this stack then move the stack to another thread and resume
+            /// it. That means that a different thread would point to our stack
+            /// and our thread doesn't point to our stack at all!
+            ///
+            /// Basically management of TLS is required here one way or another.
+            /// The strategy currently settled on is to manage the list of
+            /// activations created by this fiber as a unit. When a fiber
+            /// resumes the linked list is prepended to the current thread's
+            /// list. When the fiber is suspended then the fiber's list of
+            /// activations are all removed en-masse and saved within the fiber.
+            fn resume(&mut self, val: Result<()>) -> Result<Result<()>, ()> {
+                unsafe {
+                    let prev = self.state.take().unwrap().push();
+                    let restore = Restore {
+                        fiber: self,
+                        state: Some(prev),
+                    };
+                    return restore.fiber.fiber.resume(val);
+                }
+
+                struct Restore<'a, 'b> {
+                    fiber: &'a mut FiberFuture<'b>,
+                    state: Option<wasmtime_runtime::PreviousAsyncWasmCallState>,
+                }
+
+                impl Drop for Restore<'_, '_> {
+                    fn drop(&mut self) {
+                        unsafe {
+                            self.fiber.state = Some(self.state.take().unwrap().restore());
+                        }
+                    }
+                }
+            }
+        }
+
         impl Future for FiberFuture<'_> {
             type Output = Result<()>;
 
@@ -1694,13 +1746,34 @@ impl<T> StoreContextMut<'_, T> {
 
                     // After that's set up we resume execution of the fiber, which
                     // may also start the fiber for the first time. This either
-                    // returns `Ok` saying the fiber finished (yay!) or it returns
-                    // `Err` with the payload passed to `suspend`, which in our case
-                    // is `()`. If `Err` is returned that means the fiber polled a
-                    // future but it said "Pending", so we propagate that here.
-                    match self.fiber.resume(Ok(())) {
+                    // returns `Ok` saying the fiber finished (yay!) or it
+                    // returns `Err` with the payload passed to `suspend`, which
+                    // in our case is `()`.
+                    match self.resume(Ok(())) {
                         Ok(result) => Poll::Ready(result),
-                        Err(()) => Poll::Pending,
+
+                        // If `Err` is returned that means the fiber polled a
+                        // future but it said "Pending", so we propagate that
+                        // here.
+                        //
+                        // An additional safety check is performed when leaving
+                        // this function to help bolster the guarantees of
+                        // `unsafe impl Send` above. Notably this future may get
+                        // re-polled on a different thread. Wasmtime's
+                        // thread-local state points to the stack, however,
+                        // meaning that it would be incorrect to leave a pointer
+                        // in TLS when this function returns. This function
+                        // performs a runtime assert to verify that this is the
+                        // case, notably that the one TLS pointer Wasmtime uses
+                        // is not pointing anywhere within the stack. If it is
+                        // then that's a bug indicating that TLS management in
+                        // Wasmtime is incorrect.
+                        Err(()) => {
+                            if let Some(range) = self.fiber.stack().range() {
+                                wasmtime_runtime::AsyncWasmCallState::assert_current_state_not_in_range(range);
+                            }
+                            Poll::Pending
+                        }
                     }
                 }
             }
@@ -1724,7 +1797,7 @@ impl<T> StoreContextMut<'_, T> {
         impl Drop for FiberFuture<'_> {
             fn drop(&mut self) {
                 if !self.fiber.done() {
-                    let result = self.fiber.resume(Err(anyhow!("future dropped")));
+                    let result = self.resume(Err(anyhow!("future dropped")));
                     // This resumption with an error should always complete the
                     // fiber. While it's technically possible for host code to catch
                     // the trap and re-resume, we'd ideally like to signal that to
@@ -1732,6 +1805,8 @@ impl<T> StoreContextMut<'_, T> {
                     debug_assert!(result.is_ok());
                 }
 
+                self.state.take().unwrap().assert_null();
+
                 unsafe {
                     self.engine
                         .allocator()
@@ -1806,10 +1881,7 @@ impl AsyncCx {
                 Poll::Pending => {}
             }
 
-            let before = wasmtime_runtime::TlsRestore::take();
-            let res = (*suspend).suspend(());
-            before.replace();
-            res?;
+            (*suspend).suspend(())?;
         }
     }
 }
@@ -1945,29 +2017,31 @@ unsafe impl<T> wasmtime_runtime::Store for StoreInner<T> {
     fn new_epoch(&mut self) -> Result<u64, anyhow::Error> {
         // Temporarily take the configured behavior to avoid mutably borrowing
         // multiple times.
-        let mut behavior = std::mem::take(&mut self.epoch_deadline_behavior);
+        let mut behavior = self.epoch_deadline_behavior.take();
         let delta_result = match &mut behavior {
-            EpochDeadline::Trap => Err(Trap::Interrupt.into()),
-            EpochDeadline::Callback(callback) => {
-                let delta = callback((&mut *self).as_context_mut())?;
+            None => Err(Trap::Interrupt.into()),
+            Some(callback) => callback((&mut *self).as_context_mut()).and_then(|update| {
+                let delta = match update {
+                    UpdateDeadline::Continue(delta) => delta,
+
+                    #[cfg(feature = "async")]
+                    UpdateDeadline::Yield(delta) => {
+                        assert!(
+                            self.async_support(),
+                            "cannot use `UpdateDeadline::Yield` without enabling async support in the config"
+                        );
+                        // Do the async yield. May return a trap if future was
+                        // canceled while we're yielded.
+                        self.async_yield_impl()?;
+                        delta
+                    }
+                };
+
                 // Set a new deadline and return the new epoch deadline so
                 // the Wasm code doesn't have to reload it.
                 self.set_epoch_deadline(delta);
                 Ok(self.get_epoch_deadline())
-            }
-            #[cfg(feature = "async")]
-            EpochDeadline::YieldAndExtendDeadline { delta } => {
-                let delta = *delta;
-                // Do the async yield. May return a trap if future was
-                // canceled while we're yielded.
-                self.async_yield_impl()?;
-                // Set a new deadline.
-                self.set_epoch_deadline(delta);
-
-                // Return the new epoch deadline so the Wasm code
-                // doesn't have to reload it.
-                Ok(self.get_epoch_deadline())
-            }
+            })
         };
 
         // Put back the original behavior which was replaced by `take`.
@@ -1992,14 +2066,14 @@ impl<T> StoreInner<T> {
     }
 
     fn epoch_deadline_trap(&mut self) {
-        self.epoch_deadline_behavior = EpochDeadline::Trap;
+        self.epoch_deadline_behavior = None;
     }
 
     fn epoch_deadline_callback(
         &mut self,
-        callback: Box<dyn FnMut(StoreContextMut<T>) -> Result<u64> + Send + Sync>,
+        callback: Box<dyn FnMut(StoreContextMut<T>) -> Result<UpdateDeadline> + Send + Sync>,
     ) {
-        self.epoch_deadline_behavior = EpochDeadline::Callback(callback);
+        self.epoch_deadline_behavior = Some(callback);
     }
 
     fn epoch_deadline_async_yield_and_update(&mut self, delta: u64) {
@@ -2009,7 +2083,8 @@ impl<T> StoreInner<T> {
         );
         #[cfg(feature = "async")]
         {
-            self.epoch_deadline_behavior = EpochDeadline::YieldAndExtendDeadline { delta };
+            self.epoch_deadline_behavior =
+                Some(Box::new(move |_store| Ok(UpdateDeadline::Yield(delta))));
         }
         let _ = delta; // suppress warning in non-async build
     }
diff --git a/src/commands/run.rs b/src/commands/run.rs
index 95c5fc43e..18eec5a50 100644
--- a/src/commands/run.rs
+++ b/src/commands/run.rs
@@ -11,7 +11,7 @@ use std::thread;
 use std::time::Duration;
 use wasmtime::{
     AsContextMut, Engine, Func, GuestProfiler, Linker, Module, Store, StoreLimits,
-    StoreLimitsBuilder, Val, ValType,
+    StoreLimitsBuilder, UpdateDeadline, Val, ValType,
 };
 use wasmtime_cli_flags::{CommonOptions, WasiModules};
 use wasmtime_wasi::maybe_exit_on_error;
@@ -467,12 +467,12 @@ impl RunCommand {
                     if timeout == 0 {
                         bail!("timeout exceeded");
                     }
-                    Ok(1)
+                    Ok(UpdateDeadline::Continue(1))
                 });
             } else {
                 store.epoch_deadline_callback(move |mut store| {
                     sample(&mut store);
-                    Ok(1)
+                    Ok(UpdateDeadline::Continue(1))
                 });
             }
 
diff --git a/supply-chain/audits.toml b/supply-chain/audits.toml
index e34328225..75562d49a 100644
--- a/supply-chain/audits.toml
+++ b/supply-chain/audits.toml
@@ -721,6 +721,12 @@ are only created when requested by the application, however, and otherwise
 only does its stated purpose.
 """
 
+[[audits.file-per-thread-logger]]
+who = "Benjamin Bouvier <public@benj.me>"
+criteria = "safe-to-deploy"
+version = "0.2.0"
+notes = "Simple version bump."
+
 [[audits.file-per-thread-logger]]
 who = "Dan Gohman <dev@sunfishcode.online>"
 criteria = "safe-to-deploy"
diff --git a/winch/codegen/src/isa/aarch64/masm.rs b/winch/codegen/src/isa/aarch64/masm.rs
index e8788fe71..f092378cd 100644
--- a/winch/codegen/src/isa/aarch64/masm.rs
+++ b/winch/codegen/src/isa/aarch64/masm.rs
@@ -204,14 +204,7 @@ impl Masm for MacroAssembler {
         Address::offset(reg, offset as i64)
     }
 
-    fn cmp_with_set(
-        &mut self,
-        _dst: RegImm,
-        _lhs: RegImm,
-        _rhs: RegImm,
-        _kind: CmpKind,
-        _size: OperandSize,
-    ) {
+    fn cmp_with_set(&mut self, _src: RegImm, _dst: RegImm, _kind: CmpKind, _size: OperandSize) {
         todo!()
     }
 }
diff --git a/winch/codegen/src/isa/x64/masm.rs b/winch/codegen/src/isa/x64/masm.rs
index 0586a93ea..dfb40ae20 100644
--- a/winch/codegen/src/isa/x64/masm.rs
+++ b/winch/codegen/src/isa/x64/masm.rs
@@ -256,24 +256,9 @@ impl Masm for MacroAssembler {
         Address::offset(reg, offset)
     }
 
-    fn cmp_with_set(
-        &mut self,
-        dst: RegImm,
-        lhs: RegImm,
-        rhs: RegImm,
-        kind: CmpKind,
-        size: OperandSize,
-    ) {
-        let (src, dst): (Operand, Operand) = if dst == lhs {
-            (rhs.into(), dst.into())
-        } else {
-            panic!(
-                "the destination and first source argument must be the same, dst={:?}, lhs={:?}",
-                dst, lhs
-            );
-        };
-
-        self.asm.cmp(src, dst, size);
+    fn cmp_with_set(&mut self, src: RegImm, dst: RegImm, kind: CmpKind, size: OperandSize) {
+        let dst = dst.into();
+        self.asm.cmp(src.into(), dst, size);
         self.asm.setcc(kind, dst);
     }
 }
diff --git a/winch/codegen/src/masm.rs b/winch/codegen/src/masm.rs
index c2461b796..f5fa618b6 100644
--- a/winch/codegen/src/masm.rs
+++ b/winch/codegen/src/masm.rs
@@ -192,16 +192,9 @@ pub(crate) trait MacroAssembler {
     /// Calculate remainder.
     fn rem(&mut self, context: &mut CodeGenContext, kind: RemKind, size: OperandSize);
 
-    /// Compare two operands and put the result in dst.
+    /// Compare src and dst and put the result in dst.
     /// This function will potentially emit a series of instructions.
-    fn cmp_with_set(
-        &mut self,
-        dst: RegImm,
-        lhs: RegImm,
-        rhs: RegImm,
-        kind: CmpKind,
-        size: OperandSize,
-    );
+    fn cmp_with_set(&mut self, src: RegImm, dst: RegImm, kind: CmpKind, size: OperandSize);
 
     /// Push the register to the stack, returning the offset.
     fn push(&mut self, src: Reg) -> u32;
diff --git a/winch/codegen/src/visitor.rs b/winch/codegen/src/visitor.rs
index eaa1a7ec1..d58b1ba38 100644
--- a/winch/codegen/src/visitor.rs
+++ b/winch/codegen/src/visitor.rs
@@ -274,7 +274,7 @@ where
         use OperandSize::*;
 
         self.context.unop(self.masm, S32, &mut |masm, reg, size| {
-            masm.cmp_with_set(reg, reg, RegImm::imm(0), CmpKind::Eq, size);
+            masm.cmp_with_set(RegImm::imm(0), reg, CmpKind::Eq, size);
         });
     }
 
@@ -282,7 +282,7 @@ where
         use OperandSize::*;
 
         self.context.unop(self.masm, S64, &mut |masm, reg, size| {
-            masm.cmp_with_set(reg, reg, RegImm::imm(0), CmpKind::Eq, size);
+            masm.cmp_with_set(RegImm::imm(0), reg, CmpKind::Eq, size);
         });
     }
 
@@ -330,14 +330,14 @@ where
     fn cmp_i32s(&mut self, kind: CmpKind) {
         self.context
             .i32_binop(self.masm, &mut |masm, dst, src, size| {
-                masm.cmp_with_set(dst, dst, src, kind, size);
+                masm.cmp_with_set(src, dst, kind, size);
             });
     }
 
     fn cmp_i64s(&mut self, kind: CmpKind) {
         self.context
             .i64_binop(self.masm, &mut move |masm, dst, src, size| {
-                masm.cmp_with_set(dst, dst, src, kind, size);
+                masm.cmp_with_set(src, dst, kind, size);
             });
     }
 }


Test patch:
diff --git a/cranelift/filetests/Cargo.toml b/cranelift/filetests/Cargo.toml
index 4ff5037bd..04b4a57ed 100644
--- a/cranelift/filetests/Cargo.toml
+++ b/cranelift/filetests/Cargo.toml
@@ -18,7 +18,7 @@ cranelift-reader = { workspace = true }
 cranelift-jit = { workspace = true, features = ["selinux-fix"] }
 cranelift-module = { workspace = true }
 cranelift-control = { workspace = true }
-file-per-thread-logger = "0.1.2"
+file-per-thread-logger = { workspace = true }
 filecheck = { workspace = true }
 gimli = { workspace = true }
 log = { workspace = true }
diff --git a/tests/all/async_functions.rs b/tests/all/async_functions.rs
index a966d4677..60ec31888 100644
--- a/tests/all/async_functions.rs
+++ b/tests/all/async_functions.rs
@@ -412,16 +412,25 @@ async fn async_host_func_with_pooling_stacks() -> Result<()> {
     Ok(())
 }
 
-async fn execute_across_threads<F>(future: F) -> F::Output
+/// This will execute the `future` provided to completion and each invocation of
+/// `poll` for the future will be executed on a separate thread.
+pub async fn execute_across_threads<F>(future: F) -> F::Output
 where
     F: Future + Send + 'static,
     F::Output: Send,
 {
-    let future = PollOnce::new(Box::pin(future)).await;
-
-    tokio::task::spawn_blocking(move || tokio::runtime::Handle::current().block_on(future))
-        .await
-        .expect("shouldn't panic")
+    let mut future = Box::pin(future);
+    loop {
+        let once = PollOnce::new(future);
+        let handle = tokio::runtime::Handle::current();
+        let result = std::thread::spawn(move || handle.block_on(once))
+            .join()
+            .unwrap();
+        match result {
+            Ok(val) => break val,
+            Err(f) => future = f,
+        }
+    }
 }
 
 #[tokio::test]
@@ -696,13 +705,13 @@ impl<F> Future for PollOnce<F>
 where
     F: Future + Unpin,
 {
-    type Output = F;
+    type Output = Result<F::Output, F>;
 
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<F> {
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let mut future = self.0.take().unwrap();
         match Pin::new(&mut future).poll(cx) {
-            Poll::Pending => Poll::Ready(future),
-            Poll::Ready(_) => panic!("should not be ready"),
+            Poll::Pending => Poll::Ready(Err(future)),
+            Poll::Ready(val) => Poll::Ready(Ok(val)),
         }
     }
 }
@@ -806,7 +815,9 @@ async fn non_stacky_async_activations() -> Result<()> {
                         Ok(())
                     }
                 }) as _)
-                .await;
+                .await
+                .err()
+                .unwrap();
                 capture_stack(&stacks, &caller);
                 *caller.data_mut() = Some(future);
                 Ok(())
@@ -849,3 +860,60 @@ async fn non_stacky_async_activations() -> Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn gc_preserves_externref_on_historical_async_stacks() -> Result<()> {
+    let mut config = Config::new();
+    config.async_support(true);
+    let engine = Engine::new(&config)?;
+
+    let module = Module::new(
+        &engine,
+        r#"
+            (module $m1
+                (import "" "gc" (func $gc))
+                (import "" "recurse" (func $recurse (param i32)))
+                (import "" "test" (func $test (param i32 externref)))
+                (func (export "run") (param i32 externref)
+                    local.get 0
+                    if
+                        local.get 0
+                        i32.const -1
+                        i32.add
+                        call $recurse
+                    else
+                        call $gc
+                    end
+
+                    local.get 0
+                    local.get 1
+                    call $test
+                )
+            )
+        "#,
+    )?;
+
+    type F = TypedFunc<(i32, Option<ExternRef>), ()>;
+
+    let mut store = Store::new(&engine, None);
+    let mut linker = Linker::<Option<F>>::new(&engine);
+    linker.func_wrap("", "gc", |mut cx: Caller<'_, _>| cx.gc())?;
+    linker.func_wrap("", "test", |val: i32, handle: Option<ExternRef>| {
+        assert_eq!(handle.unwrap().data().downcast_ref(), Some(&val));
+    })?;
+    linker.func_wrap1_async("", "recurse", |mut cx: Caller<'_, _>, val: i32| {
+        let func = cx.data().unwrap();
+        Box::new(async move {
+            func.call_async(&mut cx, (val, Some(ExternRef::new(val))))
+                .await
+        })
+    })?;
+    let instance = linker.instantiate_async(&mut store, &module).await?;
+    let func: F = instance.get_typed_func(&mut store, "run")?;
+    *store.data_mut() = Some(func);
+
+    func.call_async(&mut store, (5, Some(ExternRef::new(5))))
+        .await?;
+
+    Ok(())
+}
diff --git a/tests/all/component_model/async.rs b/tests/all/component_model/async.rs
index 44877a1b0..116ced19c 100644
--- a/tests/all/component_model/async.rs
+++ b/tests/all/component_model/async.rs
@@ -1,8 +1,10 @@
 #![cfg(not(miri))]
 
+use crate::async_functions::{execute_across_threads, PollOnce};
 use anyhow::Result;
 use wasmtime::component::*;
-use wasmtime::{Store, StoreContextMut, Trap};
+use wasmtime::{Engine, Store, StoreContextMut, Trap};
+use wasmtime_component_util::REALLOC_AND_FREE;
 
 /// This is super::func::thunks, except with an async store.
 #[tokio::test]
@@ -88,3 +90,154 @@ async fn smoke_func_wrap() -> Result<()> {
 
     Ok(())
 }
+
+// This test stresses TLS management in combination with the `realloc` option
+// for imported functions. This will create an async computation which invokes a
+// component that invokes an imported function. The imported function returns a
+// list which will require invoking malloc.
+//
+// As an added stressor all polls are sprinkled across threads through
+// `execute_across_threads`. Yields are injected liberally by configuring 1
+// fuel consumption to trigger a yield.
+//
+// Overall a yield should happen during malloc which should be an "interesting
+// situation" with respect to the runtime.
+#[tokio::test]
+async fn resume_separate_thread() -> Result<()> {
+    let mut config = component_test_util::config();
+    config.async_support(true);
+    config.consume_fuel(true);
+    let engine = Engine::new(&config)?;
+    let component = format!(
+        r#"
+            (component
+                (import "yield" (func $yield (result (list u8))))
+                (core module $libc
+                    (memory (export "memory") 1)
+                    {REALLOC_AND_FREE}
+                )
+                (core instance $libc (instantiate $libc))
+
+                (core func $yield
+                    (canon lower
+                        (func $yield)
+                        (memory $libc "memory")
+                        (realloc (func $libc "realloc"))
+                    )
+                )
+
+                (core module $m
+                    (import "" "yield" (func $yield (param i32)))
+                    (import "libc" "memory" (memory 0))
+                    (func $start
+                        i32.const 8
+                        call $yield
+                    )
+                    (start $start)
+                )
+                (core instance (instantiate $m
+                    (with "" (instance (export "yield" (func $yield))))
+                    (with "libc" (instance $libc))
+                ))
+            )
+        "#
+    );
+    let component = Component::new(&engine, component)?;
+    let mut linker = Linker::new(&engine);
+    linker
+        .root()
+        .func_wrap_async("yield", |_: StoreContextMut<()>, _: ()| {
+            Box::new(async {
+                tokio::task::yield_now().await;
+                Ok((vec![1u8, 2u8],))
+            })
+        })?;
+
+    execute_across_threads(async move {
+        let mut store = Store::new(&engine, ());
+        store.out_of_fuel_async_yield(u64::MAX, 1);
+        linker.instantiate_async(&mut store, &component).await?;
+        Ok::<_, anyhow::Error>(())
+    })
+    .await?;
+    Ok(())
+}
+
+// This test is intended to stress TLS management in the component model around
+// the management of the `realloc` function. This creates an async computation
+// representing the execution of a component model function where entry into the
+// component uses `realloc` and then the component runs. This async computation
+// is then polled iteratively with another "wasm activation" (in this case a
+// core wasm function) on the stack. The poll-per-call should work and nothing
+// should in theory have problems here.
+//
+// As an added stressor all polls are sprinkled across threads through
+// `execute_across_threads`. Yields are injected liberally by configuring 1
+// fuel consumption to trigger a yield.
+//
+// Overall a yield should happen during malloc which should be an "interesting
+// situation" with respect to the runtime.
+#[tokio::test]
+async fn poll_through_wasm_activation() -> Result<()> {
+    let mut config = component_test_util::config();
+    config.async_support(true);
+    config.consume_fuel(true);
+    let engine = Engine::new(&config)?;
+    let component = format!(
+        r#"
+            (component
+                (core module $m
+                    {REALLOC_AND_FREE}
+                    (memory (export "memory") 1)
+                    (func (export "run") (param i32 i32)
+                    )
+                )
+                (core instance $i (instantiate $m))
+                (func (export "run") (param "x" (list u8))
+                    (canon lift (core func $i "run")
+                                (memory $i "memory")
+                                (realloc (func $i "realloc"))))
+            )
+        "#
+    );
+    let component = Component::new(&engine, component)?;
+    let linker = Linker::new(&engine);
+
+    let invoke_component = {
+        let engine = engine.clone();
+        async move {
+            let mut store = Store::new(&engine, ());
+            store.out_of_fuel_async_yield(u64::MAX, 1);
+            let instance = linker.instantiate_async(&mut store, &component).await?;
+            let func = instance.get_typed_func::<(Vec<u8>,), ()>(&mut store, "run")?;
+            func.call_async(&mut store, (vec![1, 2, 3],)).await?;
+            Ok::<_, anyhow::Error>(())
+        }
+    };
+
+    execute_across_threads(async move {
+        let mut store = Store::new(&engine, Some(Box::pin(invoke_component)));
+        let poll_once = wasmtime::Func::wrap0_async(&mut store, |mut cx| {
+            let invoke_component = cx.data_mut().take().unwrap();
+            Box::new(async move {
+                match PollOnce::new(invoke_component).await {
+                    Ok(result) => {
+                        result?;
+                        Ok(1)
+                    }
+                    Err(future) => {
+                        *cx.data_mut() = Some(future);
+                        Ok(0)
+                    }
+                }
+            })
+        });
+        let poll_once = poll_once.typed::<(), i32>(&mut store)?;
+        while poll_once.call_async(&mut store, ()).await? != 1 {
+            // loop around to call again
+        }
+        Ok::<_, anyhow::Error>(())
+    })
+    .await?;
+    Ok(())
+}
diff --git a/tests/all/epoch_interruption.rs b/tests/all/epoch_interruption.rs
index 8e32881c2..51d4b831c 100644
--- a/tests/all/epoch_interruption.rs
+++ b/tests/all/epoch_interruption.rs
@@ -34,7 +34,7 @@ fn make_env<T>(engine: &Engine) -> Linker<T> {
 
 enum InterruptMode {
     Trap,
-    Callback(fn(StoreContextMut<usize>) -> Result<u64>),
+    Callback(fn(StoreContextMut<usize>) -> Result<UpdateDeadline>),
     Yield(u64),
 }
 
@@ -339,7 +339,32 @@ async fn epoch_callback_continue() {
             InterruptMode::Callback(|mut cx| {
                 let s = cx.data_mut();
                 *s += 1;
-                Ok(1)
+                Ok(UpdateDeadline::Continue(1))
+            }),
+            |_| {},
+        )
+        .await
+    );
+}
+
+#[tokio::test]
+async fn epoch_callback_yield() {
+    assert_eq!(
+        Some((1, 1)),
+        run_and_count_yields_or_trap(
+            "
+            (module
+                (import \"\" \"bump_epoch\" (func $bump))
+                (func (export \"run\")
+                    call $bump  ;; bump epoch
+                    call $subfunc) ;; call func; will notice new epoch and yield
+                (func $subfunc))
+            ",
+            1,
+            InterruptMode::Callback(|mut cx| {
+                let s = cx.data_mut();
+                *s += 1;
+                Ok(UpdateDeadline::Yield(1))
             }),
             |_| {},
         )
@@ -419,7 +444,7 @@ async fn drop_future_on_epoch_yield() {
 
     let instance = linker.instantiate_async(&mut store, &module).await.unwrap();
     let f = instance.get_func(&mut store, "run").unwrap();
-    PollOnce::new(Box::pin(f.call_async(&mut store, &[], &mut []))).await;
+    let _ = PollOnce::new(Box::pin(f.call_async(&mut store, &[], &mut []))).await;
 
     assert_eq!(true, alive_flag.load(Ordering::Acquire));
 }