Evolve the wasip3 async C ABI for tasks

crates/guest-rust/src/rt/async_support/cabi.rs

@@ -43,6 +43,29 @@
 //!
 //! Additionally for now this file is serving as documentation of this
 //! interface.
+//!
+//! # Revisions
+//!
+//! This interface is intended to be evolvable over time if needed. Notably the
+//! original task structure, `wasip3_task`, has a `version` field where certain
+//! version levels imply the existence of certain fields. The historical
+//! revisions are:
+//!
+//! ### V1
+//!
+//! This was the original version. This is the original specification of
+//! `wasip3_task_set` and `wasip3_task`.
+//!
+//! ### V2
+//!
+//! This was added 2026-06-17 in response to #1618. This added
+//! `wasip3_task_v2` and `wasip3_task_vtable`. This version enables cloning a
+//! task to create a strong reference to it independent of the stack lifetime
+//! that `wasip3_task_set` is required to uphold. This necessitated introducing
+//! `clone` and `drop` callbacks to manage the lifetime of this reference.
+//! While doing this everything was moved into a vtable structure instead of
+//! inline in `wasip3_task` to make it easier to add more function pointers
+//! in the future if necessary.
 
 use core::ffi::c_void;
 
@@ -66,6 +89,7 @@ extern_wasm! {
 /// The first version of `wasip3_task` which implies the existence of the
 /// fields `ptr`, `waitable_register`, and `waitable_unregister`.
 pub const WASIP3_TASK_V1: u32 = 1;
+pub const WASIP3_TASK_V2: u32 = 2;
 
 /// Indirect "vtable" used to connect imported functions and exported tasks.
 /// Executors (e.g. exported functions) define and manage this while imports
@@ -80,6 +104,41 @@ pub struct wasip3_task {
     /// below as the first argument.
     pub ptr: *mut c_void,
 
+    /// See `wasip3_task_vtable::waitable_register`.
+    pub waitable_register: unsafe extern "C" fn(
+        ptr: *mut c_void,
+        waitable: u32,
+        callback: unsafe extern "C" fn(callback_ptr: *mut c_void, code: u32),
+        callback_ptr: *mut c_void,
+    ) -> *mut c_void,
+
+    /// See `wasip3_task_vtable::waitable_unregister`.
+    pub waitable_unregister: unsafe extern "C" fn(ptr: *mut c_void, waitable: u32) -> *mut c_void,
+}
+
+unsafe impl Send for wasip3_task {}
+unsafe impl Sync for wasip3_task {}
+
+/// Representation when `wasip3_task::version` is `WASIP3_TASK_V2`.
+#[repr(C)]
+pub struct wasip3_task_v2 {
+    /// The original task structure.
+    pub v1: wasip3_task,
+
+    /// An always-valid pointer to a list of function pointers, described
+    /// below.
+    pub vtable: &'static wasip3_task_vtable,
+}
+
+/// Function pointer operations that can operate on `wasip3_task::ptr`.
+///
+/// This was introduced in the "v2" ABI and is a member of `wasip3_task_v2`.
+#[repr(C)]
+pub struct wasip3_task_vtable {
+    /// Currently `WASIP3_TASK_V2` as that was the first version that specified
+    /// vtables.
+    pub version: u32,
+
     /// Register a new `waitable` for this exported task.
     ///
     /// This exported task will add `waitable` to its `waitable-set`. When it
@@ -104,4 +163,18 @@ pub struct wasip3_task {
     /// Returns the `callback_ptr` passed to `waitable_register` if present, or
     /// `NULL` if it's not present.
     pub waitable_unregister: unsafe extern "C" fn(ptr: *mut c_void, waitable: u32) -> *mut c_void,
+
+    /// Clones this task's pointer to create a separately owned pointer which
+    /// can be persisted outside the stack frame that this is being used
+    /// within.
+    ///
+    /// Cloned values must be dropped/deallocated with `drop` below.
+    pub clone: unsafe extern "C" fn(ptr: *mut c_void) -> *mut c_void,
+
+    /// Drops and deallocates the provided pointer previously created by a
+    /// call to the `clone` callback above.
+    ///
+    /// This must not be called on the `ptr` value within `wasip3_task::ptr` as
+    /// that's not managed with this lifetime.
+    pub drop: unsafe extern "C" fn(ptr: *mut c_void),
 }

crates/guest-rust/src/rt/async_support/inter_task_wakeup.rs

@@ -1,6 +1,6 @@
-use super::FutureState;
+use super::TaskState;
 use crate::rt::async_support::try_lock::TryLock;
-use crate::rt::async_support::{BLOCKED, COMPLETED};
+use crate::rt::async_support::{BLOCKED, COMPLETED, WaitableSet};
 use crate::{RawStreamReader, RawStreamWriter, StreamOps, UnitStreamOps};
 use core::ptr;
 
@@ -18,7 +18,7 @@ pub struct State {
     stream_reading: bool,
 }
 
-impl FutureState<'_> {
+impl TaskState<'_> {
     pub(super) fn read_inter_task_stream(&mut self) {
         // Lazily allocate the inter-task stream now that we're actually going
         // to sleep. We don't know where the wakeup notification will come from
@@ -27,7 +27,7 @@ impl FutureState<'_> {
             assert!(!self.inter_task_wakeup.stream_reading);
             let (writer, reader) = UnitStreamOps::new();
             self.inter_task_wakeup.stream = Some(reader);
-            let mut waker_stream = self.waker.inter_task_stream.lock.try_lock().unwrap();
+            let mut waker_stream = self.shared.inter_task_stream.lock.try_lock().unwrap();
             assert!(waker_stream.is_none());
             *waker_stream = Some(writer);
         }
@@ -44,7 +44,7 @@ impl FutureState<'_> {
             let rc = unsafe { UnitStreamOps.start_read(handle, ptr::null_mut(), 1) };
             assert_eq!(rc, BLOCKED);
             self.inter_task_wakeup.stream_reading = true;
-            self.add_waitable(handle);
+            self.shared.add_waitable(handle);
         }
     }
 
@@ -63,7 +63,7 @@ impl FutureState<'_> {
         unsafe {
             UnitStreamOps.cancel_read(handle);
         }
-        self.remove_waitable(handle);
+        WaitableSet::remove_waitable_from_all_sets(handle);
     }
 }
 

crates/guest-rust/src/rt/async_support/inter_task_wakeup_disabled.rs

@@ -1,9 +1,9 @@
-use super::FutureState;
+use super::TaskState;
 
 #[derive(Default)]
 pub struct State;
 
-impl FutureState<'_> {
+impl TaskState<'_> {
     pub(super) fn read_inter_task_stream(&mut self) {
         assert!(
             self.remaining_work(),