Skip to content

C/C++ API Epoch Deadline Callback #6277

@theothergraham

Description

@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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions