From a6925c21c5c4453c5f05cca92159ef9c1bd273a6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sa=C3=BAl=20Cabrera?= <saulecabrera@gmail.com>
Date: Thu, 23 Mar 2023 15:39:36 +0100
Subject: [PATCH] wasmtime: Make `StoreContextMut` accessible in epoch deadline
 callback (#6075)

This commit changes the signature of the `Store::epoch_deadline_callback` to
take in `StoreContextMut` instead of a mutable reference to the store's data.

This is useful in cases in which the callback definition needs access to the
Store to be able to use other methods that take in `AsContext`/`AsContextMut`,
like for example `WasmtimeBacktrace::capture`
---
 crates/wasmtime/src/store.rs    | 19 ++++++++++++++-----
 tests/all/epoch_interruption.rs |  5 +++--
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/crates/wasmtime/src/store.rs b/crates/wasmtime/src/store.rs
index 9a322fd6a5..8b7f5f3e61 100644
--- a/crates/wasmtime/src/store.rs
+++ b/crates/wasmtime/src/store.rs
@@ -425,11 +425,13 @@ 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(&mut T) -> Result<u64> + Send + Sync>),
+    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")]
@@ -932,7 +934,7 @@ impl<T> Store<T> {
     /// for an introduction to epoch-based interruption.
     pub fn epoch_deadline_callback(
         &mut self,
-        callback: impl FnMut(&mut T) -> Result<u64> + Send + Sync + 'static,
+        callback: impl FnMut(StoreContextMut<T>) -> Result<u64> + Send + Sync + 'static,
     ) {
         self.inner.epoch_deadline_callback(Box::new(callback));
     }
@@ -1975,10 +1977,13 @@ unsafe impl<T> wasmtime_runtime::Store for StoreInner<T> {
     }
 
     fn new_epoch(&mut self) -> Result<u64, anyhow::Error> {
-        return match &mut self.epoch_deadline_behavior {
+        // Temporarily take the configured behavior to avoid mutably borrowing
+        // multiple times.
+        let mut behavior = std::mem::take(&mut self.epoch_deadline_behavior);
+        let delta_result = match &mut behavior {
             EpochDeadline::Trap => Err(Trap::Interrupt.into()),
             EpochDeadline::Callback(callback) => {
-                let delta = callback(&mut self.data)?;
+                let delta = callback((&mut *self).as_context_mut())?;
                 // 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);
@@ -1998,6 +2003,10 @@ unsafe impl<T> wasmtime_runtime::Store for StoreInner<T> {
                 Ok(self.get_epoch_deadline())
             }
         };
+
+        // Put back the original behavior which was replaced by `take`.
+        self.epoch_deadline_behavior = behavior;
+        delta_result
     }
 }
 
@@ -2022,7 +2031,7 @@ impl<T> StoreInner<T> {
 
     fn epoch_deadline_callback(
         &mut self,
-        callback: Box<dyn FnMut(&mut T) -> Result<u64> + Send + Sync>,
+        callback: Box<dyn FnMut(StoreContextMut<T>) -> Result<u64> + Send + Sync>,
     ) {
         self.epoch_deadline_behavior = EpochDeadline::Callback(callback);
     }
diff --git a/tests/all/epoch_interruption.rs b/tests/all/epoch_interruption.rs
index c34a8b7f75..a340ef6583 100644
--- a/tests/all/epoch_interruption.rs
+++ b/tests/all/epoch_interruption.rs
@@ -32,7 +32,7 @@ fn make_env<T>(engine: &Engine) -> Linker<T> {
 
 enum InterruptMode {
     Trap,
-    Callback(fn(&mut usize) -> Result<u64>),
+    Callback(fn(StoreContextMut<usize>) -> Result<u64>),
     Yield(u64),
 }
 
@@ -334,7 +334,8 @@ async fn epoch_callback_continue() {
                 (func $subfunc))
             ",
             1,
-            InterruptMode::Callback(|s| {
+            InterruptMode::Callback(|mut cx| {
+                let s = cx.data_mut();
                 *s += 1;
                 Ok(1)
             }),