C/C++ API Epoch Deadline Callback

[theothergraham]

C/C++ API Epoch Deadline Callback

The Rust API provides pub fn epoch_deadline_callback() to add callback for the store to invoke when its epoch deadline has been exceeded, allowing the provided function to decide if the WebAssembly function should be interrupted or have its epoch deadline extended. This improvement proposes to make this functionality available in the C and C++ APIs. This is related to #3111 which proposes similar functionality for fuel exhaustion.

Benefit

This makes existing Rust API functionality available when embedding Wasmtime in C/C++.

Implementation

The below patch applies to v8.0.0. It wraps the existing Rust function with a C function that takes a C callback and void pointer, which it places in a closure. It surrounds the invocation of the callback with saving and restoring the Wasmtime runtime's TLS state, in case the function does any context switching.

vagrant@vagrant:~/edge-functions/wasmtime$ git diff
diff --git a/Cargo.lock b/Cargo.lock
index 9811cb8cb..29f22396b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3748,6 +3748,7 @@ dependencies = [
  "wasi-common",
  "wasmtime",
  "wasmtime-c-api-macros",
+ "wasmtime-runtime",
  "wasmtime-wasi",
  "wat",
 ]
diff --git a/crates/c-api/Cargo.toml b/crates/c-api/Cargo.toml
index a464c0dbd..ed860e804 100644
--- a/crates/c-api/Cargo.toml
+++ b/crates/c-api/Cargo.toml
@@ -21,6 +21,7 @@ env_logger = { workspace = true }
 anyhow = { workspace = true }
 once_cell = { workspace = true }
 wasmtime = { workspace = true, features = ['cranelift'] }
+wasmtime-runtime = { workspace = true }
 wasmtime-c-api-macros = { path = "macros" }

 # Optional dependency for the `wat2wasm` API
diff --git a/crates/c-api/include/wasmtime/store.h b/crates/c-api/include/wasmtime/store.h
index ba1d74a94..4127989c4 100644
--- a/crates/c-api/include/wasmtime/store.h
+++ b/crates/c-api/include/wasmtime/store.h
@@ -206,15 +206,29 @@ WASM_API_EXTERN wasmtime_error_t *wasmtime_context_set_wasi(wasmtime_context_t *

 /**
  * \brief Configures the relative deadline at which point WebAssembly code will
- * trap.
+ * trap or invoke the callback function.
  *
  * This function configures the store-local epoch deadline after which point
- * WebAssembly code will trap.
+ * WebAssembly code will trap or invoke the callback function.
  *
- * See also #wasmtime_config_epoch_interruption_set.
+ * See also #wasmtime_config_epoch_interruption_set and
+ * #wasmtime_store_epoch_deadline_callback.
  */
 WASM_API_EXTERN void wasmtime_context_set_epoch_deadline(wasmtime_context_t *context, uint64_t ticks_beyond_current);

+/**
+ * \brief Configures epoch deadline callback to C function.
+ *
+ * This function configures a store-local callback function that will be
+ * called when the running WebAssembly function has exceeded its epoch
+ * deadline. That function can return a 0 to raise a trap, or a greater
+ * value to add to the current epoch and resume execution of the function.
+ *
+ * See also #wasmtime_config_epoch_interruption_set and
+ * #wasmtime_context_set_epoch_deadline.
+ */
+WASM_API_EXTERN void wasmtime_store_epoch_deadline_callback(wasmtime_store_t *store, uint64_t (*func)(void*), void *data);
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif
diff --git a/crates/c-api/src/store.rs b/crates/c-api/src/store.rs
index 3949d46b0..ba8105caf 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,
+    Trap, Val,
 };

 /// This representation of a `Store` is used to implement the `wasm.h` API.
@@ -106,6 +106,44 @@ pub extern "C" fn wasmtime_store_new(
     })
 }

+// Internal structure to add Send/Sync to the c_void member.
+#[derive(Debug)]
+pub struct CallbackDataPtr {
+    pub ptr: *mut c_void,
+}
+
+impl CallbackDataPtr {
+    fn as_mut_ptr(&self) -> *mut c_void {
+        self.ptr
+    }
+}
+
+unsafe impl Send for CallbackDataPtr {}
+unsafe impl Sync for CallbackDataPtr {}
+
+// Accepts a C function pointer and opaque data pointer to invoke it with.
+// Wraps those so we can invoke the C callback via the Rust callback, and
+// surrounds the invocation with calls to save and restore the TLS state
+// in case the function is doing context switching.
+#[no_mangle]
+pub extern "C" fn wasmtime_store_epoch_deadline_callback(
+    store: &mut wasmtime_store_t,
+    func: extern "C" fn(*mut c_void) -> u64,
+    data: *mut c_void,
+) {
+    let sendable = CallbackDataPtr { ptr: data };
+    store.store.epoch_deadline_callback(move |_| {
+        let my_tls = unsafe { wasmtime_runtime::TlsRestore::take() };
+        let result = (func)(sendable.as_mut_ptr());
+        unsafe { my_tls.replace() };
+        if result > 0 {
+            Ok(result)
+        } else {
+            Err(Trap::Interrupt.into())
+        }
+    });
+}
+
 #[no_mangle]
 pub extern "C" fn wasmtime_store_context(store: &mut wasmtime_store_t) -> CStoreContextMut<'_> {
     store.store.as_context_mut()

Alternatives

I am a relative amateur with both Rust and Wasmtime, so there may be better ways or places to implement this, but this is the cleanest/smallest I could come up with.

[jameysharp]

This seems like a good idea to me!

I'd want @alexcrichton to take a look at the implementation details. Among the things I'm uncertain about:

• It looks to me like this would be the first callback exposed from the C API. I'm guessing there are going to be subtleties there.
• Should an epoch callback be able to return 0? This implementation prohibits that, using zero instead to signal that the guest should trap.

[theothergraham]

I had the same question about returning zero. After playing with it a little, I found that returning 0 is acceptable, but as far as I can tell, not necessarily useful: It will set the new deadline to the current epoch, meaning that the function will immediately trigger the callback again.

If there is some useful case for accepting 0, we can modify the C interface to return a trap/update indicator and accept a pointer to a mutable delta value for the update case.

[alexcrichton]

Agreed this is a good idea! Some thoughts I would have are:

• The callback paramter can, and probably should, take a wasmtime_context_t *ctx parameter (sourced from the closure argument provided to the callback)
• As for this being a callback, I also paused for a bit at this but wasmtime_func_new is another major way of creating callbacks so I don't think there's much more new to handle here.
• This frobs the TLS business, but I don't think that should be done here. It's not done around host functions, for example. It's only required if the whole store moves to a different thread but I'd prefer to not open that can of worms here (unless needed).
• Ideally I think the C API would mirror the Rust API, perhaps taking an out-param for the u64 return and returning an error as the actual return value. There's currently no way to programmatically create errors in the C API, though, but that also wouldn't be too hard to add. Mirroring Rust would also remove the need for reinterpreting 0 or assigning special meaning to it.

[jameysharp]

Given Alex's positive feedback, I'm looking forward to seeing a pull request for this. If you need any help following Alex's suggestions, let us know! Feel free to ask questions in this issue, or over at https://bytecodealliance.zulipchat.com/ if you prefer.

[theothergraham]

Thank you both for the fast feedback. I am working on implementing some of the changes and will get a pull request created so we can iterate on it further there.

As for the TLS business, I had a feeling that this is not the right place to do it, but I do need it. Like in #3111, I have an existing single-threaded event-based server that I need to incorporate running of multiple concurrent Wasmtime functions, and prevent any one of them from blocking the thread for extended periods of time. What I have presented in the above patch works, but perhaps we should expose TlsRestore in the C API so it can be explicitly used when it is needed.

[alexcrichton]

Ah ok, yes for something like that the need for TLS frobbing does indeed arise. Some questions for you though:

• If your server is single-threaded, why does TLS need to be saved/restored? Presumably it's not actually single-threaded, so how are wasm instances being migrated around?
• What do you do for host functions today? Do host functions never force a context switch or change of threads?
• Are you allocating stacks? If so, how?

We've been hesitant to export details of TLS business as it's not easy to get right. Even the Rust API doesn't support this (the wasmtime crate is the "official" API, what you've used here is an internal-only API which theoretically only the wasmtime crate should use). Wasmtime is built to enable this sort of cooperative switching since that's what it does for the async support in Rust, but not really any thought has been given of how to expose this to the C API. Ideally supporting suspending/etc in the C API would be done in a "first class" way than as a one-off feature of this API if we can, but I realize it's a bit of an ask too.

[theothergraham]

If your server is single-threaded, why does TLS need to be saved/restored? Presumably it's not actually single-threaded, so how are wasm instances being migrated around?

It really is single threaded, aside from a separate thread to periodically increment the epoch. Wasm instances are not being moved from thread to thread, they are being time-sliced using getcontext()/makecontext()/swapcontext() during the epoch deadline callback.

What do you do for host functions today? Do host functions never force a context switch or change of threads?

I plan to do context switching in host calls as well, as they may kick off asynchronous jobs that will need to complete (or time out) before continuing the wasm function. Today, however, I do not have this implemented.

Are you allocating stacks? If so, how?

Yes, using malloc() and placing in ucontext_t.uc_stack.ss_sp.

I am not familiar enough with the internals of TlsRestore/CallThreadState/VMRuntimeLimits to speak intelligently to why it is necessary to save and restore as I switch user-space contexts on the same thread, though it makes sense to me that it would be necessary given that async Wasmtime does it when switching contexts. I can say that without doing so, I experienced a number of crashes, stack smashing, and exiting wasm function execution on the wrong stack. Putting in the take() and replace() calls immediately cleared that up.

[alexcrichton]

Ah sorry I forget now what I was thinking, but yes even without multiple threads you still need to swap the TLS state since that's expected to be continuous per-wasm invocation. In any case definitely makes sense that this TLS swapping is required for your use case, even with one thread, and why it would remove crashes.

Hm ok though if it's desired to go down the route of "this is officially supported" then I think there's two things we could do in the C API:

• One is what you've done here, "just swap the TLS" and force embedders to take care of stack allocation, stack switching, etc. This is flexible, but it requires a lot of embedding work to get this working. For example your stacks if allocated with malloc don't have guard pages. Additionally the size of the stack isn't guaranteed to correspond with the stack size for wasm, so misconfiguration there may be possible.

• An alternative, however, is to take a more "futures" style approach that Rust uses. This would have stack switching and stack allocation all handled by Wasmtime itself (using the existing support). The complexity with this is that host functions would now have to have an interface for returning something along the lines of "please suspend me and resume later". This would look more-or-less like a coroutine.

While I realize the first is simpler and easier to implement, I would personally lean towards the second being the "official" way to support this. That meshes better I think with Wasmtime's existing support for async and stack switching, and I'd prefer to keep details like the TLS internal to Wasmtime if we can to avoid exposing too many implementation details.

[theothergraham]

Your point about my stacks not having guard pages is well taken. For my proof-of-concept that is fine, but assuming that I continue down this path, I will plan to add guard pages and tests to validate that they do their job.

As for the officially supported implementation, I'd rather take the "futures" style approach, but I have relatively little time to work on it, and my newness to Rust continues to make everything take much, much longer than I envision. So for now, I could submit a pull request for just adding a C callback for the epoch deadline, leaving out any stack switching. Obviously it will not fulfill everything I need for my use case, but stack switching should probably be its own issue/PR anyway. Does that sound good?

[alexcrichton]

Sounds reasonable to me! I can try to take a stab at a future-style API in the C API in the future perhaps too

[rockwotj]

I'm not sure if this issue is "fixed" after #6359 but related is my comment here: #3111 (comment)