wasi_hyperium/
poll.rs

1use std::{
2    collections::HashMap,
3    sync::{Arc, Mutex, Weak},
4    task::{Context, Poll, Wake, Waker},
5};
6
7use wasi::io::poll::Pollable;
8
9/// A PollableRegistry manages the polling of Pollables in relation to some
10/// Rust async executor. This must be a cheaply-`clone`able handle to its
11/// underlying state.
12pub trait PollableRegistry: Clone + Unpin {
13    type RegisteredPollable: Unpin;
14
15    /// Registers the given pollable to be polled. When the pollable is ready
16    /// the the given context's waker should be called. The pollable must be
17    /// immediately dropped when the returned RegisteredPollable is dropped.
18    fn register_pollable(&self, cx: &mut Context, pollable: Pollable) -> Self::RegisteredPollable;
19
20    /// Poll all pollables. Returns false if there are no active pollables.
21    fn poll(&self) -> bool;
22
23    /// Runs the given future to completion, polling any WASI pollables that
24    /// are registered with this registry. Returns Err(Stalled) if there are no
25    /// active pollables while the future is pending.
26    fn block_on<T>(&self, fut: impl std::future::Future<Output = T>) -> Result<T, Stalled> {
27        let mut fut = std::pin::pin!(fut);
28        let waker = noop_waker();
29        let mut cx = Context::from_waker(&waker);
30        loop {
31            if let Poll::Ready(val) = fut.as_mut().poll(&mut cx) {
32                return Ok(val);
33            }
34            if !self.poll() {
35                return Err(Stalled);
36            }
37        }
38    }
39}
40
41#[derive(Default)]
42pub struct Poller {
43    entries: Arc<Mutex<HashMap<u32, Entry>>>,
44}
45
46struct Entry {
47    pollable: Weak<Pollable>,
48    waker: Waker,
49}
50
51impl PollableRegistry for Poller {
52    type RegisteredPollable = Arc<Pollable>;
53
54    fn register_pollable(&self, cx: &mut Context, pollable: Pollable) -> Self::RegisteredPollable {
55        let handle = pollable.handle();
56        let pollable = Arc::new(pollable);
57        let entry = Entry {
58            pollable: Arc::downgrade(&pollable),
59            waker: cx.waker().clone(),
60        };
61        self.entries.lock().unwrap().insert(handle, entry);
62        pollable
63    }
64
65    fn poll(&self) -> bool {
66        let mut entries = self.entries.lock().unwrap();
67
68        // Remove any dropped pollables
69        entries.retain(|_, entry| entry.pollable.strong_count() > 0);
70
71        if entries.is_empty() {
72            return false;
73        }
74
75        // Poll pollables
76        let pollables = entries
77            .values()
78            .filter_map(|entry| entry.pollable.upgrade())
79            .collect::<Vec<_>>();
80        let pollable_refs = pollables.iter().map(|p| p.as_ref()).collect::<Vec<_>>();
81        let ready_idxs = wasi::io::poll::poll(&pollable_refs);
82
83        // Remove and wake any ready pollables
84        for idx in ready_idxs {
85            let idx: usize = idx.try_into().unwrap();
86            let handle = pollables[idx].handle();
87            let entry = entries.remove(&handle).unwrap();
88            entry.waker.wake();
89        }
90        true
91    }
92}
93
94impl Clone for Poller {
95    fn clone(&self) -> Self {
96        Self {
97            entries: self.entries.clone(),
98        }
99    }
100}
101
102pub fn noop_waker() -> Waker {
103    struct NoopWaker;
104    impl Wake for NoopWaker {
105        fn wake(self: Arc<Self>) {}
106    }
107    Arc::new(NoopWaker).into()
108}
109
110#[derive(Debug)]
111pub struct Stalled;
112
113impl std::fmt::Display for Stalled {
114    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
115        write!(f, "future stalled with no remaining pollables")
116    }
117}
118
119impl std::error::Error for Stalled {}
120
121pub trait WasiSubscribe: Unpin {
122    fn subscribe(&self) -> wasi::io::poll::Pollable;
123}
124
125macro_rules! impl_subscribe {
126    ($($ty:ty),+) => {
127        $(
128            impl WasiSubscribe for $ty {
129                fn subscribe(&self) -> wasi::io::poll::Pollable {
130                    self.subscribe()
131                }
132            }
133        )+
134    }
135}
136mod subscribe_impls {
137    use super::WasiSubscribe;
138    use wasi::http::types::*;
139    impl_subscribe!(
140        FutureTrailers,
141        InputStream,
142        OutputStream,
143        FutureIncomingResponse
144    );
145}
wasi_hyperium/poll.rs

wasi_hyperium/
poll.rs
