Implement a C ABI for async import/export communication

.github/actions/install-wasi-sdk/action.yml

@@ -5,18 +5,18 @@ runs:
   using: composite
   steps:
     - run: |
-        curl https://github.com/WebAssembly/wasi-sdk/releases/download/wasi-sdk-24/wasi-sdk-24.0-x86_64-linux.tar.gz -L | tar xzvf -
-        echo "WASI_SDK_PATH=`pwd`/wasi-sdk-24.0-x86_64-linux" >> $GITHUB_ENV
+        curl https://github.com/WebAssembly/wasi-sdk/releases/download/wasi-sdk-25/wasi-sdk-25.0-x86_64-linux.tar.gz -L | tar xzvf -
+        echo "WASI_SDK_PATH=`pwd`/wasi-sdk-25.0-x86_64-linux" >> $GITHUB_ENV
       if: runner.os == 'Linux'
       shell: bash
     - run: |
-        curl https://github.com/WebAssembly/wasi-sdk/releases/download/wasi-sdk-24/wasi-sdk-24.0-x86_64-macos.tar.gz -L | tar xzvf -
-        echo "WASI_SDK_PATH=`pwd`/wasi-sdk-24.0-x86_64-macos" >> $GITHUB_ENV
+        curl https://github.com/WebAssembly/wasi-sdk/releases/download/wasi-sdk-25/wasi-sdk-25.0-x86_64-macos.tar.gz -L | tar xzvf -
+        echo "WASI_SDK_PATH=`pwd`/wasi-sdk-25.0-x86_64-macos" >> $GITHUB_ENV
       if: runner.os == 'macOS'
       shell: bash
     - run: |
-        curl https://github.com/WebAssembly/wasi-sdk/releases/download/wasi-sdk-24/wasi-sdk-24.0-x86_64-windows.tar.gz -L | tar xzvf -
-        echo "WASI_SDK_PATH=`pwd`/wasi-sdk-24.0-x86_64-windows" >> $GITHUB_ENV
+        curl https://github.com/WebAssembly/wasi-sdk/releases/download/wasi-sdk-25/wasi-sdk-25.0-x86_64-windows.tar.gz -L | tar xzvf -
+        echo "WASI_SDK_PATH=`pwd`/wasi-sdk-25.0-x86_64-windows" >> $GITHUB_ENV
       if: runner.os == 'Windows'
       shell: bash
     - name: Setup `wasm-tools`

.github/workflows/main.yml

@@ -151,10 +151,10 @@ jobs:
     - name: Install wasm32-wasip1 target
       run: rustup target add wasm32-wasip1
 
-    # Verify the output of the `./ci/rebuild-libcabi-realloc.sh` script is
+    # Verify the output of the `./ci/rebuild-libwit-bindgen-cabi.sh` script is
     # up-to-date.
     - uses: ./.github/actions/install-wasi-sdk
-    - run: ./ci/rebuild-libcabi-realloc.sh
+    - run: ./ci/rebuild-libwit-bindgen-cabi.sh
     - run: git diff --exit-code
 
     # Test various feature combinations, make sure they all build

ci/rebuild-libcabi-realloc.sh → ci/rebuild-libwit-bindgen-cabi.sh

@@ -41,13 +41,16 @@ set -ex
 
 version=$(./ci/print-current-version.sh | sed 's/\./_/g')
 
-sym=cabi_realloc_wit_bindgen_$version
+realloc=cabi_realloc_wit_bindgen_$version
 
-cat >./crates/guest-rust/rt/src/cabi_realloc.rs <<-EOF
+rm -f crates/guest-rust/rt/src/wit_bindgen_*.{rs,o,c}
+rm -f crates/guest-rust/rt/src/libwit_bindgen_cabi.a
+
+cat >./crates/guest-rust/rt/src/wit_bindgen_cabi_realloc.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,
@@ -57,29 +60,48 @@ pub unsafe extern "C" fn $sym(
 }
 EOF
 
-cat >./crates/guest-rust/rt/src/cabi_realloc.c <<-EOF
+cat >./crates/guest-rust/rt/src/wit_bindgen_cabi_realloc.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);
 }
 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
+cat >./crates/guest-rust/rt/src/wit_bindgen_cabi_wasip3.c <<-EOF
+// This file is generated by $0
+
+#include <stdlib.h>
+
+static void *WASIP3_TASK = NULL;
+
+__attribute__((__weak__))
+void *wasip3_task_set(void *ptr) {
+  void *ret = WASIP3_TASK;
+  WASIP3_TASK = ptr;
+  return ret;
+}
+EOF
+
+build() {
+  file=$1
+  $WASI_SDK_PATH/bin/clang crates/guest-rust/rt/src/$1.c \
+    -O -c -o crates/guest-rust/rt/src/$1.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/$1.o \
+    -o ./crates/guest-rust/rt/src/$1.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
+build wit_bindgen_cabi_realloc
+build wit_bindgen_cabi_wasip3
 
-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
+$WASI_SDK_PATH/bin/llvm-ar crus crates/guest-rust/rt/src/libwit_bindgen_cabi.a \
+  crates/guest-rust/rt/src/wit_bindgen_cabi_realloc.o \
+  crates/guest-rust/rt/src/wit_bindgen_cabi_wasip3.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(),