Implement a C ABI for async import/export communication

This commit is an implementation of a solution for
WebAssembly/component-model#485 for Rust. This should enable releasing
multiple versions of `wit-bindgen` into the wild and have them all work
together for now. Integration with `wasi-libc` will come in the future
in theory.

Author: alexcrichton

ci/rebuild-libcabi-realloc.sh ci/rebuild-libwit-bindgen-cabi.shci/rebuild-libcabi-realloc.sh renamed to ci/rebuild-libwit-bindgen-cabi.sh

@@ -41,45 +41,65 @@ set -ex
 
 version=$(./ci/print-current-version.sh | sed 's/\./_/g')
 
-sym=cabi_realloc_wit_bindgen_$version
+realloc=cabi_realloc_wit_bindgen_$version
+wasip3_task_set=wasip3_task_set_wit_bindgen_$version
+wasip3_task_get=wasip3_task_get_wit_bindgen_$version
 
-cat >./crates/guest-rust/rt/src/cabi_realloc.rs <<-EOF
+cat >./crates/guest-rust/rt/src/wit_bindgen_cabi.rs <<-EOF
 // This file is generated by $0
 
 #[unsafe(no_mangle)]
-pub unsafe extern "C" fn $sym(
+pub unsafe extern "C" fn $realloc(
     old_ptr: *mut u8,
     old_len: usize,
     align: usize,
     new_len: usize,
 ) -> *mut u8 {
     crate::cabi_realloc(old_ptr, old_len, align, new_len)
 }
+
+static mut WASIP3_TASK: *mut u8 = core::ptr::null_mut();
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn $wasip3_task_set(ptr: *mut u8) -> *mut u8 {
+    unsafe {
+        let ret = WASIP3_TASK;
+        WASIP3_TASK = ptr;
+        ret
+    }
+}
 EOF
 
-cat >./crates/guest-rust/rt/src/cabi_realloc.c <<-EOF
+cat >./crates/guest-rust/rt/src/wit_bindgen_cabi.c <<-EOF
 // This file is generated by $0
 
 #include <stdint.h>
 
-extern void *$sym(void *ptr, size_t old_size, size_t align, size_t new_size);
+extern void *$realloc(void *ptr, size_t old_size, size_t align, size_t new_size);
 
 __attribute__((__weak__, __export_name__("cabi_realloc")))
 void *cabi_realloc(void *ptr, size_t old_size, size_t align, size_t new_size) {
-  return $sym(ptr, old_size, align, new_size);
+  return $realloc(ptr, old_size, align, new_size);
+}
+
+extern void *$wasip3_task_set(void *ptr);
+
+__attribute__((__weak__))
+void *wasip3_task_set(void *ptr) {
+  return $wasip3_task_set(ptr);
 }
 EOF
 
-rm -f crates/guest-rust/rt/src/cabi_realloc.o
-$WASI_SDK_PATH/bin/clang crates/guest-rust/rt/src/cabi_realloc.c \
-  -O -c -o crates/guest-rust/rt/src/cabi_realloc.o
+rm -f crates/guest-rust/rt/src/wit_bindgen_cabi.o
+$WASI_SDK_PATH/bin/clang crates/guest-rust/rt/src/wit_bindgen_cabi.c \
+  -O -c -o crates/guest-rust/rt/src/wit_bindgen_cabi.o
 
 # Remove the `producers` section. This appears to differ whether the host for
 # clang is either macOS or Linux. Not needed here anyway, so discard it to help
 # either host produce the same object.
-wasm-tools strip -d producers ./crates/guest-rust/rt/src/cabi_realloc.o \
-  -o ./crates/guest-rust/rt/src/cabi_realloc.o
+wasm-tools strip -d producers ./crates/guest-rust/rt/src/wit_bindgen_cabi.o \
+  -o ./crates/guest-rust/rt/src/wit_bindgen_cabi.o
 
-rm -f crates/guest-rust/rt/src/libwit_bindgen_cabi_realloc.a
-$WASI_SDK_PATH/bin/llvm-ar crus crates/guest-rust/rt/src/libwit_bindgen_cabi_realloc.a \
-  crates/guest-rust/rt/src/cabi_realloc.o
+rm -f crates/guest-rust/rt/src/libwit_bindgen_cabi.a
+$WASI_SDK_PATH/bin/llvm-ar crus crates/guest-rust/rt/src/libwit_bindgen_cabi.a \
+  crates/guest-rust/rt/src/wit_bindgen_cabi.o

crates/guest-rust/rt/README.md

@@ -8,13 +8,13 @@ crate is to contain "runtime" code related to the macro-expansion of the
 be removed in some situations.
 
 This crate contains a precompiled object file and archive at
-`src/cabi_realloc.o` and `src/libwit_bindgen_cabi_realloc.a`. This is compiled
-from the source `src/cabi_realloc.c` and is checked in as precompiled to avoid
-needing a C compiler at compile-time which isn't always available. This object
-file is only used on wasm targets.
+`src/wit_bindgen_cabi.o` and `src/libwit_bindgen_cabi.a`. This is compiled
+from the source `src/wit_bindgen_cabi.c` and is checked in as precompiled to
+avoid needing a C compiler at compile-time which isn't always available. This
+object file is only used on wasm targets.
 
 The object file is compiled by
-[this script]https://github.com/bytecodealliance/wit-bindgen/blob/main/ci/rebuild-libcabi-realloc.sh)
+[this script]https://github.com/bytecodealliance/wit-bindgen/blob/main/ci/rebuild-libwit-bindgen-cabi.sh)
 and is verified in repository continuous integration that the checked-in
 versions match what CI produces.
 

crates/guest-rust/rt/build.rs

@@ -16,10 +16,10 @@ fn main() {
 
     let mut src = env::current_dir().unwrap();
     src.push("src");
-    src.push("libwit_bindgen_cabi_realloc.a");
+    src.push("libwit_bindgen_cabi.a");
 
     let dst_name = format!(
-        "wit_bindgen_cabi_realloc{}",
+        "wit_bindgen_cabi{}",
         env!("CARGO_PKG_VERSION").replace(".", "_")
     );
     let dst = out_dir.join(format!("lib{dst_name}.a"));

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

@@ -8,14 +8,15 @@ extern crate std;
 use core::sync::atomic::{AtomicBool, Ordering};
 use std::boxed::Box;
 use std::collections::HashMap;
+use std::ffi::c_void;
 use std::fmt::{self, Debug, Display};
 use std::future::Future;
 use std::mem;
 use std::pin::Pin;
 use std::ptr;
 use std::string::String;
 use std::sync::Arc;
-use std::task::{Context, Poll, Wake, Waker};
+use std::task::{Context, Poll, Wake};
 use std::vec::Vec;
 
 use futures::channel::oneshot;
@@ -36,6 +37,7 @@ macro_rules! rtdebug {
 }
 
 mod abi_buffer;
+mod cabi;
 mod future_support;
 mod stream_support;
 mod waitable;
@@ -60,21 +62,32 @@ struct FutureState {
     tasks: Option<FuturesUnordered<BoxFuture>>,
     /// The waitable set containing waitables created by this task, if any.
     waitable_set: Option<u32>,
-    /// A map of waitables to the corresponding waker and completion code.
-    ///
-    /// This is primarily filled in and managed by `WaitableOperation<S>`. The
-    /// waker here comes straight from `std::task::Context` and the pointer is
-    /// otherwise stored within the `WaitableOperation<S>` The raw pointer here
-    /// has a disconnected lifetime with each future but the management of the
-    /// internal states with respect to drop should always ensure that this is
-    /// only ever pointing to active waitable operations.
-    ///
-    /// When a waitable notification is received the corresponding entry in this
-    /// map is removed, the status code is filled in, and the waker is notified.
-    wakers: HashMap<u32, (Waker, *mut Option<u32>)>,
+
+    /// State of all waitables in `waitable_set`, and the ptr/callback they're
+    /// associated with.
+    waitables: HashMap<u32, (*mut c_void, unsafe extern "C" fn(*mut c_void, u32))>,
+
+    /// Raw structure used to pass to `cabi::wasip3_task_set`
+    wasip3_task: cabi::wasip3_task,
 }
 
 impl FutureState {
+    fn new(future: BoxFuture) -> FutureState {
+        FutureState {
+            todo: 0,
+            tasks: Some([future].into_iter().collect()),
+            waitable_set: None,
+            waitables: HashMap::new(),
+            wasip3_task: cabi::wasip3_task {
+                // This pointer is filled in before calling `wasip3_task_set`.
+                ptr: ptr::null_mut(),
+                version: cabi::WASIP3_TASK_V1,
+                waitable_register,
+                waitable_unregister,
+            },
+        }
+    }
+
     fn get_or_create_waitable_set(&mut self) -> u32 {
         *self.waitable_set.get_or_insert_with(waitable_set_new)
     }
@@ -88,7 +101,32 @@ impl FutureState {
     }
 
     fn remaining_work(&self) -> bool {
-        self.todo > 0 || !self.wakers.is_empty()
+        self.todo > 0 || !self.waitables.is_empty()
+    }
+}
+
+unsafe extern "C" fn waitable_register(
+    ptr: *mut c_void,
+    waitable: u32,
+    callback: unsafe extern "C" fn(*mut c_void, u32),
+    callback_ptr: *mut c_void,
+) -> *mut c_void {
+    let ptr = ptr.cast::<FutureState>();
+    assert!(!ptr.is_null());
+    (*ptr).add_waitable(waitable);
+    match (*ptr).waitables.insert(waitable, (callback_ptr, callback)) {
+        Some((prev, _)) => prev,
+        None => ptr::null_mut(),
+    }
+}
+
+unsafe extern "C" fn waitable_unregister(ptr: *mut c_void, waitable: u32) -> *mut c_void {
+    let ptr = ptr.cast::<FutureState>();
+    assert!(!ptr.is_null());
+    (*ptr).remove_waitable(waitable);
+    match (*ptr).waitables.remove(&waitable) {
+        Some((prev, _)) => prev,
+        None => ptr::null_mut(),
     }
 }
 
@@ -145,6 +183,22 @@ unsafe fn poll(state: *mut FutureState) -> Poll<()> {
         }
     }
 
+    // Finish our `wasip3_task` by initializing its self-referential pointer,
+    // and then register it for the duration of this function with
+    // `wasip3_task_set`. The previous value of `wasip3_task_set` will get
+    // restored when this function returns.
+    struct ResetTask(*mut cabi::wasip3_task);
+    impl Drop for ResetTask {
+        fn drop(&mut self) {
+            unsafe {
+                cabi::wasip3_task_set(self.0);
+            }
+        }
+    }
+    (*state).wasip3_task.ptr = state.cast();
+    let prev = cabi::wasip3_task_set(&mut (*state).wasip3_task);
+    let _reset = ResetTask(prev);
+
     loop {
         if let Some(futures) = (*state).tasks.as_mut() {
             let old = CURRENT;
@@ -191,16 +245,9 @@ pub fn first_poll<T: 'static>(
     future: impl Future<Output = T> + 'static,
     fun: impl FnOnce(&T) + 'static,
 ) -> i32 {
-    let state = Box::into_raw(Box::new(FutureState {
-        todo: 0,
-        tasks: Some(
-            [Box::pin(future.map(|v| fun(&v))) as BoxFuture]
-                .into_iter()
-                .collect(),
-        ),
-        waitable_set: None,
-        wakers: HashMap::new(),
-    }));
+    let state = Box::into_raw(Box::new(FutureState::new(Box::pin(
+        future.map(|v| fun(&v)),
+    ))));
     let done = unsafe { poll(state).is_ready() };
     unsafe { callback_code(state, done) }
 }
@@ -339,9 +386,8 @@ unsafe fn callback_with_state(
                 "EVENT_{{STREAM,FUTURE}}_{{READ,WRITE}}({event0:#x}, {event1:#x}, {event2:#x})"
             );
             (*state).remove_waitable(event1 as u32);
-            let (waker, code) = (*state).wakers.remove(&(event1 as u32)).unwrap();
-            *code = Some(event2 as u32);
-            waker.wake();
+            let (ptr, callback) = (*state).waitables.remove(&(event1 as u32)).unwrap();
+            callback(ptr, event2 as u32);
 
             let done = poll(state).is_ready();
             callback_code(state, done)
@@ -469,16 +515,7 @@ pub fn spawn(future: impl Future<Output = ()> + 'static) {
 // TODO: refactor so `'static` bounds aren't necessary
 pub fn block_on<T: 'static>(future: impl Future<Output = T> + 'static) -> T {
     let (tx, mut rx) = oneshot::channel();
-    let state = &mut FutureState {
-        todo: 0,
-        tasks: Some(
-            [Box::pin(future.map(move |v| drop(tx.send(v)))) as BoxFuture]
-                .into_iter()
-                .collect(),
-        ),
-        waitable_set: None,
-        wakers: HashMap::new(),
-    };
+    let state = &mut FutureState::new(Box::pin(future.map(move |v| drop(tx.send(v)))) as BoxFuture);
     loop {
         match unsafe { poll(state) } {
             Poll::Ready(()) => break rx.try_recv().unwrap().unwrap(),

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

@@ -0,0 +1,105 @@
+//! Definition of the "C ABI" of how imported functions interact with exported
+//! tasks.
+//!
+//! Ok this crate is written in Rust, why in the world does this exist? This
+//! comment is intended to explain this rationale but the tl;dr; is we want
+//! this to work:
+//!
+//! * Within a single component ...
+//! * One rust crate uses `wit-bindgen 0.A.0` to generate an exported function.
+//! * One rust crate uses `wit-bindgen 0.B.0` to bind an imported function.
+//! * The two crates are connected in the application with
+//!   `std::future::Future`.
+//!
+//! Without this module this situation won't work because 0.A.0 has no
+//! knowledge of 0.B.0 meaning that when 0.B.0 decides to block it won't know
+//! where to register its `waitable` within a `waitable-set`.
+//!
+//! To solve this problem the long-term intention is that something will live
+//! in `wasi-libc` itself, but in the meantime it's living "somewhere" within
+//! `wit-bindgen 0.*.0`. Specifically all `wit-bindgen` versions will all
+//! reference, via C linkage, a single function which is used to manipulate a
+//! single pointer in linear memory. This pointer is a `wasip3_task` structure
+//! which has all the various fields to use it.
+//!
+//! The `wasip3_task_set` symbol is itself defined in C inside of the
+//! `src/wit_bindgen_cabi.c` file at this time, specifically because it's
+//! annotated with `__weak__` meaning that any definition of it suffices. This
+//! isn't possible to define in stable Rust (specifically `__weak__`).
+//!
+//! Once `wasip3_task_set` is defined everything then operates via indirection,
+//! aka based off the returned pointer. The intention is that exported functions
+//! will set this (it's sort of like an executor) and then imported functions
+//! will all use this as the source of registering waitables. In the end that
+//! means that it's possible to share types with `std::future::Future` that
+//! are backed at the ABI level with this "channel".
+//!
+//! In the future it's hoped that this can move into `wasi-libc` itself, or if
+//! `wasi-libc` provides something else that would be prioritized over this.
+//! For now this is basically an affordance that we're going to be frequently
+//! releaseing new major versions of `wit-bindgen` and we don't want to force
+//! applications to all be using the exact same version of the bindings
+//! generator and async bindings.
+//!
+//! Additionally for now this file is serving as documentation of this
+//! interface.
+
+use core::ffi::c_void;
+
+extern "C" {
+    /// Sets the global task pointer to `ptr` provided. Returns the previous
+    /// value.
+    ///
+    /// This function acts as both a dual getter and a setter. To get the
+    /// current task pointer a dummy `ptr` can be provided (e.g. NULL) and then
+    /// it's passed back when you're done working with it. When setting the
+    /// current task pointer it's recommended to call this and then call it
+    /// again with the previous value when the tasks's work is done.
+    ///
+    /// For executors they need to ensure that the `ptr` passed in lives for
+    /// the entire lifetime of the component model task.
+    pub fn wasip3_task_set(ptr: *mut wasip3_task) -> *mut wasip3_task;
+}
+
+/// 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;
+
+/// Indirect "vtable" used to connect imported functions and exported tasks.
+/// Executors (e.g. exported functions) define and manage this while imports
+/// use it.
+#[repr(C)]
+pub struct wasip3_task {
+    /// Currently `WASIP3_TASK_V1`. Indicates what fields are present next
+    /// depending on the version here.
+    pub version: u32,
+
+    /// Private pointer owned by the `wasip3_task` itself, passed to callbacks
+    /// below as the first argument.
+    pub ptr: *mut c_void,
+
+    /// Register a new `waitable` for this exported task.
+    ///
+    /// This exported task will add `waitable` to its `waitable-set`. When it
+    /// becomes ready then `callback` will be invoked with the ready code as
+    /// well as the `callback_ptr` provided.
+    ///
+    /// If `waitable` was previously registered with this task then the
+    /// previuos `callback_ptr` is returned. Otherwise `NULL` is returned.
+    ///
+    /// It's the caller's responsibility to ensure that `callback_ptr` is valid
+    /// until `callback` is invoked, `waitable_unregister` is invoked, or
+    /// `waitable_register` is called again to overwrite the value.
+    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,
+
+    /// Removes the `waitable` from this task's `waitable-set`.
+    ///
+    /// 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,
+}