spin_executor/
lib.rs

1use bindings::wasi::io;
2use std::future::Future;
3use std::mem;
4use std::sync::{Arc, Mutex};
5use std::task::{Context, Poll, Wake, Waker};
6
7/// Module containing the generated WIT bindings.
8pub mod bindings {
9    wit_bindgen::generate!({
10        world: "imports",
11        path: "io.wit",
12    });
13}
14
15impl std::fmt::Display for io::streams::Error {
16    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
17        f.write_str(&self.to_debug_string())
18    }
19}
20
21impl std::error::Error for io::streams::Error {}
22
23static WAKERS: Mutex<Vec<(io::poll::Pollable, Waker)>> = Mutex::new(Vec::new());
24
25/// Push a Pollable and Waker to WAKERS.
26pub fn push_waker(pollable: io::poll::Pollable, waker: Waker) {
27    WAKERS.lock().unwrap().push((pollable, waker));
28}
29
30/// Run the specified future to completion blocking until it yields a result.
31///
32/// Based on an executor using `wasi::io/poll/poll-list`,
33pub fn run<T>(future: impl Future<Output = T>) -> T {
34    futures::pin_mut!(future);
35    struct DummyWaker;
36
37    impl Wake for DummyWaker {
38        fn wake(self: Arc<Self>) {}
39    }
40
41    let waker = Arc::new(DummyWaker).into();
42
43    loop {
44        match future.as_mut().poll(&mut Context::from_waker(&waker)) {
45            Poll::Pending => {
46                let mut new_wakers = Vec::new();
47
48                let wakers = mem::take::<Vec<_>>(&mut WAKERS.lock().unwrap());
49
50                assert!(!wakers.is_empty());
51
52                let pollables = wakers
53                    .iter()
54                    .map(|(pollable, _)| pollable)
55                    .collect::<Vec<_>>();
56
57                let mut ready = vec![false; wakers.len()];
58
59                for index in io::poll::poll(&pollables) {
60                    ready[usize::try_from(index).unwrap()] = true;
61                }
62
63                for (ready, (pollable, waker)) in ready.into_iter().zip(wakers) {
64                    if ready {
65                        waker.wake()
66                    } else {
67                        new_wakers.push((pollable, waker));
68                    }
69                }
70
71                *WAKERS.lock().unwrap() = new_wakers;
72            }
73            Poll::Ready(result) => break result,
74        }
75    }
76}
spin_executor/lib.rs

spin_executor/
lib.rs
