c-api: support yielding in wasmtime_store_epoch_deadline_callback

crates/c-api/include/wasmtime/store.h

@@ -200,19 +200,44 @@ WASM_API_EXTERN wasmtime_error_t *wasmtime_context_set_wasi(wasmtime_context_t *
  */
 WASM_API_EXTERN void wasmtime_context_set_epoch_deadline(wasmtime_context_t *context, uint64_t ticks_beyond_current);
 
+/// \brief An enum for the behavior before extending the epoch deadline.
+typedef uint8_t wasmtime_update_deadline_kind_t;
+/// \brief Directly continue to updating the deadline and executing WebAssembly.
+#define WASMTIME_UPDATE_DEADLINE_CONTINUE 0
+/// \brief Yield control (via async support) then update the deadline.
+#define WASMTIME_UPDATE_DEADLINE_YIELD 1
+
 /**
  * \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 #wasmtime_error_t to terminate
- * the function, or set the delta argument and return NULL to update the
- * epoch deadline and resume function execution.
+ * deadline. That function can:
+ * - return a #wasmtime_error_t to terminate the function
+ * - set the delta argument and return NULL to update the
+ * epoch deadline delta and resume function execution.
+ * - set the delta argument, update the epoch deadline, 
+ * set update_kind to WASMTIME_UPDATE_DEADLINE_YIELD,
+ * and return NULL to yield (via async support) and 
+ * resume function execution.
+ *
+ * To use WASMTIME_UPDATE_DEADLINE_YIELD async support must be enabled 
+ * for this store.
  *
  * 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, wasmtime_error_t* (*func)(wasmtime_context_t*, void*, uint64_t*), void *data);
+WASM_API_EXTERN void wasmtime_store_epoch_deadline_callback(
+ wasmtime_store_t *store, 
+ wasmtime_error_t* (*func)(
+ wasmtime_context_t* context, 
+ void* data, 
+ uint64_t* epoch_deadline_delta, 
+ wasmtime_update_deadline_kind_t* update_kind
+ ), 
+ void *data,
+ void (*finalizer)(void*)
+);
 
 #ifdef __cplusplus
 } // extern "C"

crates/c-api/src/async.rs

@@ -87,16 +87,22 @@ impl Future for wasmtime_async_continuation_t {
  }
 }
 
-pub type wasmtime_func_async_continuation_callback_t = extern "C" fn(*mut c_void) -> bool;
-
-struct CallbackData {
- env: *mut c_void,
+/// Internal structure to add Send/Sync to a c_void member.
+///
+/// This is useful in closures that need to capture some C data.
+#[derive(Debug)]
+struct CallbackDataPtr {
+ pub ptr: *mut std::ffi::c_void,
 }
-unsafe impl Send for CallbackData {}
+
+unsafe impl Send for CallbackDataPtr {}
+unsafe impl Sync for CallbackDataPtr {}
+
+pub type wasmtime_func_async_continuation_callback_t = extern "C" fn(*mut c_void) -> bool;
 
 async fn invoke_c_async_callback<'a>(
  cb: wasmtime_func_async_callback_t,
- data: CallbackData,
+ data: CallbackDataPtr,
  mut caller: Caller<'a, crate::StoreData>,
  params: &'a [Val],
  results: &'a mut [Val],
@@ -127,7 +133,7 @@ async fn invoke_c_async_callback<'a>(
  finalizer: None,
  };
  cb(
- data.env,
+ data.ptr,
  &mut caller,
  params.as_ptr(),
  params.len(),
@@ -171,7 +177,7 @@ unsafe fn c_async_callback_to_rust_fn(
  let foreign = crate::ForeignData { data, finalizer };
  move |caller, params, results| {
  let _ = &foreign; // move entire foreign into this closure
- let data = CallbackData { env: foreign.data };
+ let data = CallbackDataPtr { ptr: foreign.data };
  Box::new(invoke_c_async_callback(
  callback, data, caller, params, results,
  ))

crates/c-api/src/store.rs

@@ -106,20 +106,9 @@ 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 {}
+pub type wasmtime_update_deadline_kind_t = u8;
+pub const WASMTIME_UPDATE_DEADLINE_CONTINUE: wasmtime_update_deadline_kind_t = 0;
+pub const WASMTIME_UPDATE_DEADLINE_YIELD: wasmtime_update_deadline_kind_t = 1;
 
 #[no_mangle]
 pub extern "C" fn wasmtime_store_epoch_deadline_callback(
@@ -128,22 +117,32 @@ pub extern "C" fn wasmtime_store_epoch_deadline_callback(
  CStoreContextMut<'_>,
  *mut c_void,
  *mut u64,
+ *mut wasmtime_update_deadline_kind_t,
  ) -> Option<Box<wasmtime_error_t>>,
  data: *mut c_void,
+ finalizer: Option<extern "C" fn(*mut c_void)>,
 ) {
- let sendable = CallbackDataPtr { ptr: data };
+ let foreign = crate::ForeignData { data, finalizer };
  store.store.epoch_deadline_callback(move |mut store_ctx| {
+ let _ = &foreign; // Move foreign into this closure
  let mut delta: u64 = 0;
+ let mut kind = WASMTIME_UPDATE_DEADLINE_CONTINUE;
  let result = (func)(
  store_ctx.as_context_mut(),
- sendable.as_mut_ptr(),
+ foreign.data,
  &mut delta as *mut u64,
+ &mut kind as *mut wasmtime_update_deadline_kind_t,
  );
  match result {
  Some(err) => Err(wasmtime::Error::from(<wasmtime_error_t as Into<
  anyhow::Error,
  >>::into(*err))),
- None => Ok(UpdateDeadline::Continue(delta)),
+ None if kind == WASMTIME_UPDATE_DEADLINE_CONTINUE => {
+ Ok(UpdateDeadline::Continue(delta))
+ }
+ #[cfg(feature = "async")]
+ None if kind == WASMTIME_UPDATE_DEADLINE_YIELD => Ok(UpdateDeadline::Yield(delta)),
+ _ => panic!("unknown wasmtime_update_deadline_kind_t: {}", kind),
  }
  });
 }