wit_bindgen/rt/async_support/
spawn.rs

1// TODO: Switch to interior mutability (e.g. use Mutexes or thread-local
2// RefCells) and remove this, since even in single-threaded mode `static mut`
3// references can be a hazard due to recursive access.
4#![allow(static_mut_refs)]
5
6use crate::rt::async_support::BoxFuture;
7use alloc::boxed::Box;
8use alloc::vec::Vec;
9use core::future::Future;
10use core::task::{Context, Poll};
11use futures::stream::{FuturesUnordered, StreamExt};
12
13/// Any newly-deferred work queued by calls to the `spawn` function while
14/// polling the current task.
15static mut SPAWNED: Vec<BoxFuture> = Vec::new();
16
17#[derive(Default)]
18pub struct Tasks<'a> {
19    tasks: FuturesUnordered<BoxFuture<'a>>,
20}
21
22impl<'a> Tasks<'a> {
23    pub fn new(root: BoxFuture<'a>) -> Tasks<'a> {
24        Tasks {
25            tasks: [root].into_iter().collect(),
26        }
27    }
28
29    pub fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll<()> {
30        loop {
31            // Perform some work by seeing what's next in this
32            // `FuturesUnordered`. Afterwards check the set of spawned tasks
33            // and, if any, add them to our set of tasks being done.
34            let poll = self.tasks.poll_next_unpin(cx);
35            let spawned = unsafe {
36                if SPAWNED.is_empty() {
37                    false
38                } else {
39                    self.tasks.extend(SPAWNED.drain(..));
40                    true
41                }
42            };
43            match poll {
44                // If no tasks were ready, and if we didn't spawn any work,
45                // then there's nothing left to do so return pending.
46                //
47                // If no tasks were ready, and if we spawned some work, then
48                // turn the loop again to register interest in the work and
49                // ensure that it's not forgotten about.
50                Poll::Pending => {
51                    if !spawned {
52                        return Poll::Pending;
53                    }
54                }
55
56                // If our set of tasks is empty it shouldn't be possible to have
57                // spawned anything, and return saying that we're done.
58                Poll::Ready(None) => {
59                    assert!(!spawned);
60                    return Poll::Ready(());
61                }
62
63                // If a task finished, then turn the loop again to see if there
64                // are any other completed tasks. This also serves double-duty
65                // to ensure that we look at everything in our set of tasks
66                // before concluding that we're finished.
67                Poll::Ready(Some(())) => {}
68            }
69        }
70    }
71
72    pub fn is_empty(&self) -> bool {
73        self.tasks.is_empty()
74    }
75}
76
77/// Spawn the provided `future` to get executed concurrently with the
78/// currently-running async computation.
79///
80/// This API is somewhat similar to `tokio::task::spawn` for example but has a
81/// number of limitations to be aware of. If possible it's recommended to avoid
82/// this, but it can be convenient if these limitations do not apply to you:
83///
84/// * Spawned tasks do not work when the version of the `wit-bindgen` crate
85///   managing the export bindings is different from the version of this crate.
86///   To work correctly the `spawn` function and export executor must be at
87///   exactly the same version. Given the major-version-breaking nature of
88///   `wit-bindgen` this is not always easy to rely on. This is tracked in
89///   [#1305].
90///
91/// * Spawned tasks do not outlive the scope of the async computation they are
92///   spawned within. For example with an async export function spawned tasks
93///   will be polled within the context of that component-model async task. For
94///   computations executing within a [`block_on`] call, however, the spawned
95///   tasks will be executed within that scope. This notably means that for
96///   [`block_on`] spawned tasks will prevent the [`block_on`] function from
97///   returning, even if a value is available to return.
98///
99/// * There is no handle returned to the spawned task meaning that it cannot be
100///   cancelled or monitored.
101///
102/// * The task spawned here is executed *concurrently*, not in *parallel*. This
103///   means that while one future is being polled no other future can be polled
104///   at the same time. This is similar to a single-thread executor in Tokio.
105///
106/// With these restrictions in mind this can be used to express
107/// execution-after-returning in the component model. For example once an
108/// exported async function has produced a value this can be used to continue to
109/// execute some more code before the component model async task exits.
110///
111/// [`block_on`]: crate::block_on
112/// [#1305]: https://github.com/bytecodealliance/wit-bindgen/issues/1305
113pub fn spawn_local(future: impl Future<Output = ()> + 'static) {
114    unsafe { SPAWNED.push(Box::pin(future)) }
115}
wit_bindgen/rt/async_support/spawn.rs

wit_bindgen/rt/async_support/
spawn.rs
