c-api: support yielding in epoch deadline callback

Signed-off-by: Tyler Rockwood <rockwood@redpanda.com>

Author: rockwotj

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 wasmtime_update_deadline_kind_t uint8_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/store.rs

@@ -121,29 +121,42 @@ impl CallbackDataPtr {
 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(
     store: &mut wasmtime_store_t,
     func: extern "C" fn(
         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))
+            }
+            None if kind == WASMTIME_UPDATE_DEADLINE_YIELD => Ok(UpdateDeadline::Yield(delta)),
+            _ => panic!("unknown wasmtime_update_deadline_kind_t: {}", kind),
         }
     });
 }