Skip to main content

wit_bindgen/rt/
async_support.rs

1#![deny(missing_docs)]
2
3use self::try_lock::TryLock;
4use alloc::boxed::Box;
5use alloc::collections::BTreeMap;
6use alloc::sync::Arc;
7use alloc::task::Wake;
8use core::ffi::c_void;
9use core::future::Future;
10use core::mem::{self, ManuallyDrop};
11use core::pin::Pin;
12use core::ptr;
13use core::sync::atomic::{AtomicU32, Ordering};
14use core::task::{Context, Poll, Waker};
15
16macro_rules! rtdebug {
17    ($($f:tt)*) => {
18        // Change this flag to enable debugging, right now we're not using a
19        // crate like `log` or such to reduce runtime deps. Intended to be used
20        // during development for now.
21        if false {
22            #[cfg(feature = "std")]
23            std::eprintln!($($f)*);
24        }
25    }
26}
27
28/// Helper macro to deduplicate foreign definitions of wasm functions.
29///
30/// This automatically imports when on wasm targets and then defines a dummy
31/// panicking shim for native targets to support native compilation but fail at
32/// runtime.
33macro_rules! extern_wasm {
34    (
35        $(#[$extern_attr:meta])*
36        unsafe extern "C" {
37            $(
38                $(#[$func_attr:meta])*
39                $vis:vis fn $func_name:ident ( $($args:tt)* ) $(-> $ret:ty)?;
40            )*
41        }
42    ) => {
43        $(
44            #[cfg(not(target_family = "wasm"))]
45            #[allow(unused, reason = "dummy shim for non-wasm compilation, never invoked")]
46            $vis unsafe fn $func_name($($args)*) $(-> $ret)? {
47                unreachable!();
48            }
49        )*
50
51        #[cfg(target_family = "wasm")]
52        $(#[$extern_attr])*
53        unsafe extern "C" {
54            $(
55                $(#[$func_attr])*
56                $vis fn $func_name($($args)*) $(-> $ret)?;
57            )*
58        }
59    };
60}
61
62mod abi_buffer;
63mod cabi;
64mod error_context;
65mod future_support;
66#[cfg(feature = "futures-stream")]
67mod futures_stream;
68#[cfg(feature = "inter-task-wakeup")]
69mod inter_task_wakeup;
70mod stream_support;
71mod subtask;
72mod try_lock;
73#[cfg(feature = "inter-task-wakeup")]
74mod unit_stream;
75mod waitable;
76mod waitable_set;
77
78#[cfg(not(feature = "inter-task-wakeup"))]
79use inter_task_wakeup_disabled as inter_task_wakeup;
80#[cfg(not(feature = "inter-task-wakeup"))]
81mod inter_task_wakeup_disabled;
82
83use self::waitable_set::WaitableSet;
84pub use abi_buffer::*;
85pub use error_context::*;
86pub use future_support::*;
87#[cfg(feature = "futures-stream")]
88pub use futures_stream::*;
89pub use stream_support::*;
90#[doc(hidden)]
91pub use subtask::Subtask;
92#[cfg(feature = "inter-task-wakeup")]
93pub use unit_stream::*;
94
95type BoxFuture<'a> = Pin<Box<dyn Future<Output = ()> + 'a>>;
96
97#[cfg(feature = "async-spawn")]
98mod spawn;
99#[cfg(feature = "async-spawn")]
100pub use spawn::{JoinHandle, spawn_local};
101#[cfg(not(feature = "async-spawn"))]
102mod spawn_disabled;
103#[cfg(not(feature = "async-spawn"))]
104use spawn_disabled as spawn;
105
106/// Represents a task created by either a call to an async-lifted export or a
107/// future run using `block_on` or `start_task`.
108struct TaskState<'a> {
109    /// Remaining work to do (if any) before this task can be considered "done".
110    ///
111    /// Note that we won't tell the host the task is done until this is drained
112    /// and `waitables` is empty.
113    tasks: spawn::Tasks<'a>,
114
115    /// Dual-mode rust-level Waker and C ABI level "task" for wasip3
116    /// integration.
117    shared: Arc<SharedTaskState>,
118
119    /// Clone of `shared` field, but represented as `std::task::Waker`.
120    waker: Waker,
121
122    /// State related to supporting inter-task wakeup scenarios.
123    inter_task_wakeup: inter_task_wakeup::State,
124}
125
126#[derive(Default)]
127struct SharedTaskState {
128    /// One of `SLEEP_STATE_*` indicating the current status.
129    sleep_state: AtomicU32,
130    inter_task_stream: inter_task_wakeup::WakerState,
131
132    /// State of all waitables in `waitable_set`, and the ptr/callback they're
133    /// associated with.
134    //
135    // Note that this is a `BTreeMap` rather than a `HashMap` only because, as
136    // of this writing, initializing the default hasher for `HashMap` requires
137    // calling `wasi_snapshot_preview1:random_get`, which requires initializing
138    // the `wasi_snapshot_preview1` adapter when targeting `wasm32-wasip2` and
139    // later, and that's expensive enough that we'd prefer to avoid it for apps
140    // which otherwise make no use of the adapter.
141    //
142    // Also note that the `TryLock` here should never be contended, but it's
143    // used for interior mutability.
144    waitables: TryLock<BTreeMap<u32, CabiWaitable>>,
145
146    /// The waitable set containing waitables created by this task, if any.
147    //
148    // Note the `TryLock` is the same as `waitables` above, it's serving the
149    // purpose of interior mutability.
150    waitable_set: TryLock<Option<WaitableSet>>,
151}
152
153/// An entry of `SharedTaskState::waitables` which is added through the C ABI.
154struct CabiWaitable {
155    callback: unsafe extern "C" fn(*mut c_void, u32),
156    callback_ptr: *mut c_void,
157}
158
159unsafe impl Send for CabiWaitable {}
160
161impl<'a> TaskState<'a> {
162    fn new(future: impl Future<Output = ()> + 'a) -> TaskState<'a> {
163        let shared = Arc::new(SharedTaskState::default());
164        TaskState {
165            waker: shared.clone().into(),
166            shared,
167            tasks: spawn::Tasks::new(future),
168            inter_task_wakeup: Default::default(),
169        }
170    }
171
172    fn remaining_work(&self) -> bool {
173        !self.shared.waitables.try_lock().unwrap().is_empty()
174    }
175
176    /// Handles the `event{0,1,2}` event codes and returns a corresponding
177    /// return code along with a flag whether this future is "done" or not.
178    fn callback(&mut self, event0: u32, event1: u32, event2: u32) -> CallbackCode {
179        match event0 {
180            EVENT_NONE => rtdebug!("EVENT_NONE"),
181            EVENT_SUBTASK => rtdebug!("EVENT_SUBTASK({event1:#x}, {event2:#x})"),
182            EVENT_STREAM_READ => rtdebug!("EVENT_STREAM_READ({event1:#x}, {event2:#x})"),
183            EVENT_STREAM_WRITE => rtdebug!("EVENT_STREAM_WRITE({event1:#x}, {event2:#x})"),
184            EVENT_FUTURE_READ => rtdebug!("EVENT_FUTURE_READ({event1:#x}, {event2:#x})"),
185            EVENT_FUTURE_WRITE => rtdebug!("EVENT_FUTURE_WRITE({event1:#x}, {event2:#x})"),
186
187            // Delegate cancellation to the set of spawned tasks that this
188            // component-model task is maintaining. The cancel here is a signal
189            // for the "root computation", but everything else spawned, if any,
190            // gets to continue along its merry way.
191            EVENT_CANCEL => {
192                rtdebug!("EVENT_CANCEL");
193                self.with_p3_task_set(|me| me.tasks.cancel_root());
194            }
195            _ => unreachable!(),
196        }
197
198        self.with_p3_task_set(|me| {
199            // Transition our sleep state to ensure that the inter-task stream
200            // isn't used since there's no need to use that here.
201            me.shared
202                .sleep_state
203                .store(SLEEP_STATE_WOKEN, Ordering::Relaxed);
204
205            // With all of our context now configured, deliver the event
206            // notification this callback corresponds to.
207            //
208            // Note that this should happen under the reset of
209            // `waker.sleep_state` above to ensure that if a waker is woken it
210            // won't actually signal our inter-task stream since we're already
211            // in the process of handling the future.
212            if event0 != EVENT_NONE && event0 != EVENT_CANCEL {
213                me.deliver_waitable_event(event1, event2)
214            }
215
216            // If there's still an in-progress read (e.g. `event{1,2}`) wasn't
217            // ourselves getting woken up, then cancel the read since we're
218            // processing the future here anyway.
219            me.cancel_inter_task_stream_read();
220
221            loop {
222                let mut context = Context::from_waker(&me.waker);
223
224                // On each turn of this loop reset the state to "polling"
225                // which clears out any pending wakeup if one was sent. This
226                // in theory helps minimize wakeups from previous iterations
227                // happening in this iteration.
228                me.shared
229                    .sleep_state
230                    .store(SLEEP_STATE_POLLING, Ordering::Relaxed);
231
232                // Poll our future, seeing if it was able to make progress.
233                let poll = me.tasks.poll_next(&mut context);
234
235                match poll {
236                    // The task list is empty, but there might be remaining work
237                    // in terms of waitables through the cabi interface. In this
238                    // situation wait for all waitables to be resolved before
239                    // signaling that our own task is done.
240                    Poll::Ready(()) => {
241                        assert!(me.tasks.is_empty());
242                        if me.remaining_work() {
243                            let set = me.shared.waitable_set.try_lock().unwrap();
244                            let waitable = set.as_ref().unwrap().as_raw();
245                            break CallbackCode::Wait(waitable);
246                        } else {
247                            break CallbackCode::Exit;
248                        }
249                    }
250
251                    // Some future within `self.tasks` is not ready yet. If our
252                    // `waker` was signaled then that means this is a yield
253                    // operation, otherwise it means we're blocking on
254                    // something.
255                    Poll::Pending => {
256                        assert!(!me.tasks.is_empty());
257                        if me.shared.sleep_state.load(Ordering::Relaxed) == SLEEP_STATE_WOKEN {
258                            if me.remaining_work() {
259                                let (event0, event1, event2) = {
260                                    let set = me.shared.waitable_set.try_lock().unwrap();
261                                    set.as_ref().unwrap().poll()
262                                };
263                                if event0 != EVENT_NONE {
264                                    me.deliver_waitable_event(event1, event2);
265                                    continue;
266                                }
267                            }
268                            break CallbackCode::Yield;
269                        }
270
271                        // Transition our state to "sleeping" so wakeup
272                        // notifications know that they need to signal the
273                        // inter-task stream.
274                        me.shared
275                            .sleep_state
276                            .store(SLEEP_STATE_SLEEPING, Ordering::Relaxed);
277                        me.read_inter_task_stream();
278                        let set = me.shared.waitable_set.try_lock().unwrap();
279                        let waitable = set.as_ref().unwrap().as_raw();
280                        break CallbackCode::Wait(waitable);
281                    }
282                }
283            }
284        })
285    }
286
287    /// Deliver the `code` event to the `waitable` store within our map. This
288    /// waitable should be present because it's part of the waitable set which
289    /// is kept in-sync with our map.
290    fn deliver_waitable_event(&mut self, waitable: u32, code: u32) {
291        WaitableSet::remove_waitable_from_all_sets(waitable);
292
293        if self
294            .inter_task_wakeup
295            .consume_waitable_event(waitable, code)
296        {
297            return;
298        }
299
300        let c = {
301            let mut waitables = self.shared.waitables.try_lock().unwrap();
302            waitables.remove(&waitable).unwrap()
303        };
304        unsafe {
305            (c.callback)(c.callback_ptr, code);
306        }
307    }
308
309    fn with_p3_task_set<R>(&mut self, f: impl FnOnce(&mut Self) -> R) -> R {
310        // Initialize a temporary `wasip3_task` structure on the stack and
311        // inform `wasip3_task_set` that we're now within that task. Note the
312        // RAII guard to reset the task back to its previous contents.
313        struct ResetTask(*mut cabi::wasip3_task);
314        impl Drop for ResetTask {
315            fn drop(&mut self) {
316                unsafe {
317                    cabi::wasip3_task_set(self.0);
318                }
319            }
320        }
321        // The `ptr` field of `wasip3_task` is to `SharedTaskState` which is
322        // what's cloned/handed out/etc.
323        let shared_raw: *const SharedTaskState = &*self.shared;
324        let mut wasip3_task = cabi::wasip3_task_v3 {
325            v1: cabi::wasip3_task {
326                ptr: shared_raw.cast_mut().cast(),
327                version: cabi::WASIP3_TASK_V3,
328                waitable_register: SharedTaskState::CABI_VTABLE.v2.waitable_register,
329                waitable_unregister: SharedTaskState::CABI_VTABLE.v2.waitable_unregister,
330            },
331            vtable: &SharedTaskState::CABI_VTABLE,
332        };
333
334        // Explicitly take a mutable borrow on the entire `wasip3_task`
335        // structure, and then cast its raw pointer to the "smaller" historical
336        // version, ensuring the final pointer has provenace over the entire
337        // structure.
338        let wasip3_task: *mut cabi::wasip3_task_v3 = &mut wasip3_task;
339        let prev = unsafe { cabi::wasip3_task_set(wasip3_task.cast::<cabi::wasip3_task>()) };
340        let _reset = ResetTask(prev);
341
342        f(self)
343    }
344}
345
346impl Drop for TaskState<'_> {
347    fn drop(&mut self) {
348        // No further wakeup notifications are needed once destruction starts.
349        // Destructors may wake other futures in this task, so suppress those
350        // notifications before cancelling the inter-task stream read.
351        self.shared
352            .sleep_state
353            .store(SLEEP_STATE_WOKEN, Ordering::Relaxed);
354
355        // If there's an active read of the inter-task stream, go ahead and
356        // cancel it, since we're about to drop the stream anyway.
357        self.cancel_inter_task_stream_read();
358
359        // If this state has active tasks then they need to be dropped which may
360        // execute arbitrary code. This arbitrary code might require the p3 APIs
361        // for managing waitables, notably around removing them. In this
362        // situation we ensure that the p3 task is set while futures are being
363        // destroyed.
364        if !self.tasks.is_empty() {
365            self.with_p3_task_set(|me| {
366                me.tasks = Default::default();
367            })
368        }
369    }
370}
371
372impl SharedTaskState {
373    const CABI_VTABLE: cabi::wasip3_task_vtable_v3 = cabi::wasip3_task_vtable_v3 {
374        v2: cabi::wasip3_task_vtable {
375            waitable_register: Self::cabi_waitable_register,
376            waitable_unregister: Self::cabi_waitable_unregister,
377            drop: Self::cabi_drop,
378            clone: Self::cabi_clone,
379        },
380        #[cfg(feature = "async-spawn")]
381        rust_spawn: Some(Self::cabi_rust_spawn),
382        #[cfg(not(feature = "async-spawn"))]
383        rust_spawn: None,
384    };
385
386    /// Adds the `waitable` provided to this task's waitable set.
387    fn add_waitable(&self, waitable: u32) {
388        let mut set = self.waitable_set.try_lock().unwrap();
389        set.get_or_insert_with(WaitableSet::new).join(waitable);
390    }
391
392    /// Implementation of the CABI `waitable_register` function.
393    fn waitable_register(
394        &self,
395        waitable: u32,
396        callback: unsafe extern "C" fn(*mut c_void, u32),
397        callback_ptr: *mut c_void,
398    ) -> *mut c_void {
399        self.add_waitable(waitable);
400        let mut waitables = self.waitables.try_lock().unwrap();
401        let c = CabiWaitable {
402            callback,
403            callback_ptr,
404        };
405        match waitables.insert(waitable, c) {
406            Some(prev) => prev.callback_ptr,
407            None => ptr::null_mut(),
408        }
409    }
410
411    /// Implementation of the CABI `waitable_unregister` function.
412    fn waitable_unregister(&self, waitable: u32) -> *mut c_void {
413        WaitableSet::remove_waitable_from_all_sets(waitable);
414        let mut waitables = self.waitables.try_lock().unwrap();
415        match waitables.remove(&waitable) {
416            Some(prev) => prev.callback_ptr,
417            None => ptr::null_mut(),
418        }
419    }
420
421    /// Helper to go from a raw `c_void` FFI pointer to a typed
422    /// self-representation.
423    unsafe fn cabi_to_self(ptr: *mut c_void) -> ManuallyDrop<Arc<SharedTaskState>> {
424        unsafe { ManuallyDrop::new(Arc::from_raw(ptr.cast::<SharedTaskState>())) }
425    }
426
427    unsafe extern "C" fn cabi_waitable_register(
428        ptr: *mut c_void,
429        waitable: u32,
430        callback: unsafe extern "C" fn(*mut c_void, u32),
431        callback_ptr: *mut c_void,
432    ) -> *mut c_void {
433        let me = unsafe { Self::cabi_to_self(ptr) };
434        me.waitable_register(waitable, callback, callback_ptr)
435    }
436
437    unsafe extern "C" fn cabi_waitable_unregister(ptr: *mut c_void, waitable: u32) -> *mut c_void {
438        let me = unsafe { Self::cabi_to_self(ptr) };
439        me.waitable_unregister(waitable)
440    }
441
442    unsafe extern "C" fn cabi_clone(ptr: *mut c_void) -> *mut c_void {
443        let me = unsafe { Self::cabi_to_self(ptr) };
444        Arc::into_raw(Arc::clone(&me)).cast_mut().cast()
445    }
446
447    unsafe extern "C" fn cabi_drop(ptr: *mut c_void) {
448        let mut me = unsafe { Self::cabi_to_self(ptr) };
449        unsafe { ManuallyDrop::drop(&mut me) }
450    }
451
452    #[cfg(feature = "async-spawn")]
453    unsafe fn cabi_rust_spawn(_ptr: *mut c_void, task: Box<dyn Future<Output = ()>>) {
454        spawn::push(Pin::from(task))
455    }
456}
457
458/// Status for "this task is actively being polled"
459const SLEEP_STATE_POLLING: u32 = 0;
460/// Status for "no further wakeup notification is needed", either because the
461/// task has already been woken or because it is being destroyed.
462const SLEEP_STATE_WOKEN: u32 = 1;
463/// Status for "this task is not being polled and has not been woken"
464///
465/// Wakeups on this status signal the inter-task stream.
466const SLEEP_STATE_SLEEPING: u32 = 2;
467
468impl Wake for SharedTaskState {
469    fn wake(self: Arc<Self>) {
470        Self::wake_by_ref(&self)
471    }
472
473    fn wake_by_ref(self: &Arc<Self>) {
474        match self.sleep_state.swap(SLEEP_STATE_WOKEN, Ordering::Relaxed) {
475            // If this future was currently being polled, or if someone else
476            // already woke it up, then there's nothing to do.
477            SLEEP_STATE_POLLING | SLEEP_STATE_WOKEN => {}
478
479            // If this future is sleeping, however, then this is a cross-task
480            // wakeup meaning that we need to write to its wakeup stream.
481            other => {
482                assert_eq!(other, SLEEP_STATE_SLEEPING);
483                self.inter_task_stream.wake();
484            }
485        }
486    }
487}
488
489const EVENT_NONE: u32 = 0;
490const EVENT_SUBTASK: u32 = 1;
491const EVENT_STREAM_READ: u32 = 2;
492const EVENT_STREAM_WRITE: u32 = 3;
493const EVENT_FUTURE_READ: u32 = 4;
494const EVENT_FUTURE_WRITE: u32 = 5;
495const EVENT_CANCEL: u32 = 6;
496
497#[derive(PartialEq, Debug)]
498enum CallbackCode {
499    Exit,
500    Yield,
501    Wait(u32),
502}
503
504impl CallbackCode {
505    fn encode(self) -> u32 {
506        match self {
507            CallbackCode::Exit => 0,
508            CallbackCode::Yield => 1,
509            CallbackCode::Wait(waitable) => 2 | (waitable << 4),
510        }
511    }
512}
513
514const STATUS_STARTING: u32 = 0;
515const STATUS_STARTED: u32 = 1;
516const STATUS_RETURNED: u32 = 2;
517const STATUS_STARTED_CANCELLED: u32 = 3;
518const STATUS_RETURNED_CANCELLED: u32 = 4;
519
520const BLOCKED: u32 = 0xffff_ffff;
521const COMPLETED: u32 = 0x0;
522const DROPPED: u32 = 0x1;
523const CANCELLED: u32 = 0x2;
524
525/// Return code of stream/future operations.
526#[derive(PartialEq, Debug, Copy, Clone)]
527enum ReturnCode {
528    /// The operation is blocked and has not completed.
529    Blocked,
530    /// The operation completed with the specified number of items.
531    Completed(u32),
532    /// The other end is dropped, but before that the specified number of items
533    /// were transferred.
534    Dropped(u32),
535    /// The operation was cancelled, but before that the specified number of
536    /// items were transferred.
537    Cancelled(u32),
538}
539
540impl ReturnCode {
541    fn decode(val: u32) -> ReturnCode {
542        if val == BLOCKED {
543            return ReturnCode::Blocked;
544        }
545        let amt = val >> 4;
546        match val & 0xf {
547            COMPLETED => ReturnCode::Completed(amt),
548            DROPPED => ReturnCode::Dropped(amt),
549            CANCELLED => ReturnCode::Cancelled(amt),
550            _ => panic!("unknown return code {val:#x}"),
551        }
552    }
553}
554
555#[cfg(not(target_env = "p3"))]
556mod task_state {
557    pub fn get() -> *mut u8 {
558        extern_wasm! {
559            #[link(wasm_import_module = "$root")]
560            unsafe extern "C" {
561                #[link_name = "[context-get-0]"]
562                fn get() -> *mut u8;
563            }
564        }
565
566        unsafe { get() }
567    }
568
569    pub unsafe fn set(value: *mut u8) {
570        extern_wasm! {
571            #[link(wasm_import_module = "$root")]
572            unsafe extern "C" {
573                #[link_name = "[context-set-0]"]
574                fn set(value: *mut u8);
575            }
576        }
577
578        unsafe { set(value) }
579    }
580}
581
582#[cfg(all(target_env = "p3", feature = "std"))]
583mod task_state {
584    std::thread_local!(static TASK_STATE: std::cell::Cell<*mut u8> = std::cell::Cell::new(std::ptr::null_mut()));
585
586    pub fn get() -> *mut u8 {
587        TASK_STATE.with(|c| c.get())
588    }
589
590    pub unsafe fn set(value: *mut u8) {
591        TASK_STATE.with(|c| c.set(value))
592    }
593}
594
595#[cfg(all(target_env = "p3", not(feature = "std")))]
596mod wasip3_context;
597#[cfg(all(target_env = "p3", not(feature = "std")))]
598mod task_state {
599
600    pub fn get() -> *mut u8 {
601        unsafe { super::wasip3_context::get() }
602    }
603
604    pub unsafe fn set(value: *mut u8) {
605        unsafe { super::wasip3_context::set(value) }
606    }
607}
608
609/// Starts execution of the `task` provided, an asynchronous computation.
610///
611/// This is used for async-lifted exports at their definition site. The
612/// representation of the export is `task` and this function is called from the
613/// entrypoint. The code returned here is the same as the callback associated
614/// with this export, and the callback will be used if this task doesn't exit
615/// immediately with its result.
616#[doc(hidden)]
617pub fn start_task(task: impl Future<Output = ()> + 'static) -> i32 {
618    // Allocate a new `TaskState` which will track all state necessary for
619    // our exported task.
620    let state = Box::into_raw(Box::new(TaskState::new(task)));
621
622    // Store our `TaskState` into our context-local-storage slot and then
623    // pretend we got EVENT_NONE to kick off everything.
624    //
625    // SAFETY: we should own `context.set` as we're the root level exported
626    // task, and then `callback` is only invoked when context-local storage is
627    // valid.
628    unsafe {
629        assert!(task_state::get().is_null());
630        task_state::set(state.cast());
631        callback(EVENT_NONE, 0, 0) as i32
632    }
633}
634
635/// Handle a progress notification from the host regarding either a call to an
636/// async-lowered import or a stream/future read/write operation.
637///
638/// # Unsafety
639///
640/// This function assumes that `context_get()` returns a `TaskState`.
641#[doc(hidden)]
642pub unsafe fn callback(event0: u32, event1: u32, event2: u32) -> u32 {
643    // Acquire our context-local state, assert it's not-null, and then reset
644    // the state to null while we're running to help prevent any unintended
645    // usage.
646    let state = task_state::get().cast::<TaskState<'static>>();
647    assert!(!state.is_null());
648    unsafe {
649        task_state::set(ptr::null_mut());
650    }
651
652    // Use `state` to run the `callback` function in the context of our event
653    // codes we received. If the callback decides to exit then we're done with
654    // our future so deallocate it. Otherwise put our future back in
655    // context-local storage and forward the code.
656    unsafe {
657        let rc = (*state).callback(event0, event1, event2);
658        if rc == CallbackCode::Exit {
659            drop(Box::from_raw(state));
660        } else {
661            task_state::set(state.cast());
662        }
663        rtdebug!(" => (cb) {rc:?}");
664        rc.encode()
665    }
666}
667
668/// Run the specified future to completion, returning the result.
669///
670/// This uses `waitable-set.wait` to poll for progress on any in-progress calls
671/// to async-lowered imports as necessary.
672pub fn block_on<T>(future: impl Future<Output = T>) -> T {
673    let mut result = None;
674    let mut state = TaskState::new(async {
675        result = Some(future.await);
676    });
677    let mut event = (EVENT_NONE, 0, 0);
678    loop {
679        match state.callback(event.0, event.1, event.2) {
680            CallbackCode::Exit => {
681                drop(state);
682                break result.unwrap();
683            }
684            // When yielding there's not much else to do here other than just go
685            // back and poll again -- the `callback` function will already poll
686            // the waitable set if applicable so just turn the loop and deliver
687            // an `EVENT_NONE` event.
688            CallbackCode::Yield => {
689                event = (EVENT_NONE, 0, 0);
690            }
691            CallbackCode::Wait(_) => {
692                let set = state.shared.waitable_set.try_lock().unwrap();
693                event = set.as_ref().unwrap().wait()
694            }
695        }
696    }
697}
698
699/// Call the `yield` canonical built-in function.
700///
701/// This yields control to the host temporarily, allowing other tasks to make
702/// progress. It's a good idea to call this inside a busy loop which does not
703/// otherwise ever yield control the host.
704///
705/// Note that this function is a blocking function, not an `async` function.
706/// That means that this is not an async yield which allows other tasks in this
707/// component to progress, but instead this will block the current function
708/// until the host gets back around to returning from this yield. Asynchronous
709/// functions should probably use [`yield_async`] instead.
710///
711/// # Return Value
712///
713/// This function returns a `bool` which indicates whether execution should
714/// continue after this yield point. A return value of `true` means that the
715/// task was not cancelled and execution should continue. A return value of
716/// `false`, however, means that the task was cancelled while it was suspended
717/// at this yield point. The caller should return back and exit from the task
718/// ASAP in this situation.
719pub fn yield_blocking() -> bool {
720    extern_wasm! {
721        #[link(wasm_import_module = "$root")]
722        unsafe extern "C" {
723            #[link_name = "[thread-yield]"]
724            fn yield_() -> bool;
725        }
726    }
727
728    // Note that the return value from the raw intrinsic is inverted, the
729    // canonical ABI returns "did this task get cancelled" while this function
730    // works as "should work continue going".
731    unsafe { !yield_() }
732}
733
734/// The asynchronous counterpart to [`yield_blocking`].
735///
736/// This function does not block the current task but instead gives the
737/// Rust-level executor a chance to yield control back to the host temporarily.
738/// This means that other Rust-level tasks may also be able to progress during
739/// this yield operation.
740///
741/// # Return Value
742///
743/// Unlike [`yield_blocking`] this function does not return anything. If this
744/// component task is cancelled while paused at this yield point then the future
745/// will be dropped and a Rust-level destructor will take over and clean up the
746/// task. It's not necessary to do anything with the return value of this
747/// function other than ensuring that you `.await` the function call.
748pub async fn yield_async() {
749    #[derive(Default)]
750    struct Yield {
751        yielded: bool,
752    }
753
754    impl Future for Yield {
755        type Output = ();
756
757        fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<()> {
758            if self.yielded {
759                Poll::Ready(())
760            } else {
761                self.yielded = true;
762                context.waker().wake_by_ref();
763                Poll::Pending
764            }
765        }
766    }
767
768    Yield::default().await;
769}
770
771/// Call the `backpressure.inc` canonical built-in function.
772pub fn backpressure_inc() {
773    extern_wasm! {
774        #[link(wasm_import_module = "$root")]
775        unsafe extern "C" {
776            #[link_name = "[backpressure-inc]"]
777            fn backpressure_inc();
778        }
779    }
780
781    unsafe { backpressure_inc() }
782}
783
784/// Call the `backpressure.dec` canonical built-in function.
785pub fn backpressure_dec() {
786    extern_wasm! {
787        #[link(wasm_import_module = "$root")]
788        unsafe extern "C" {
789            #[link_name = "[backpressure-dec]"]
790            fn backpressure_dec();
791        }
792    }
793
794    unsafe { backpressure_dec() }
795}
796
797#[doc(hidden)]
798pub struct TaskCancelOnDrop {
799    _priv: (),
800}
801
802impl TaskCancelOnDrop {
803    #[doc(hidden)]
804    pub fn new() -> TaskCancelOnDrop {
805        TaskCancelOnDrop { _priv: () }
806    }
807
808    #[doc(hidden)]
809    pub fn forget(self) {
810        mem::forget(self);
811    }
812}
813
814impl Drop for TaskCancelOnDrop {
815    fn drop(&mut self) {
816        extern_wasm! {
817            #[link(wasm_import_module = "[export]$root")]
818            unsafe extern "C" {
819                #[link_name = "[task-cancel]"]
820                fn cancel();
821            }
822        }
823
824        unsafe { cancel() }
825    }
826}