Skip to main content

wasmedge_sys/async/
fiber.rs

1//! Defines data structure for WasmEdge async mechanism.
2
3use fiber_for_wasmedge::{Fiber, FiberStack, Suspend};
4use std::{
5    future::Future,
6    pin::Pin,
7    ptr,
8    task::{Context, Poll},
9};
10
11/// Defines a FiberFuture.
12pub(crate) struct FiberFuture<'a> {
13    fiber: Fiber<'a, Result<(), ()>, (), Result<(), ()>>,
14    current_suspend: *mut *const Suspend<Result<(), ()>, (), Result<(), ()>>,
15    current_poll_cx: *mut *mut Context<'static>,
16}
17impl<'a> FiberFuture<'a> {
18    /// Create a fiber to execute the given function.
19    ///
20    /// # Arguments
21    ///
22    /// * `async_state` - Used to store asynchronous state at run time.
23    ///
24    /// * `func` - The function to execute.
25    ///
26    /// # Error
27    ///
28    /// If fail to create the fiber stack or the fiber fail to resume, then an error is returned.
29    pub(crate) async fn on_fiber<R>(
30        async_state: &AsyncState,
31        func: impl FnOnce() -> R + Send,
32    ) -> Result<R, ()> {
33        let mut slot = None;
34        let future = {
35            let current_poll_cx = async_state.current_poll_cx.get();
36            let current_suspend = async_state.current_suspend.get();
37
38            let stack = FiberStack::new(2 << 20).map_err(|_e| ())?;
39            let slot = &mut slot;
40            let fiber = Fiber::new(stack, move |keep_going, suspend| {
41                keep_going?;
42
43                unsafe {
44                    let _reset = Reset(current_suspend, *current_suspend);
45                    *current_suspend = suspend;
46                    *slot = Some(func());
47                    Ok(())
48                }
49            })
50            .map_err(|_e| ())?;
51
52            FiberFuture {
53                fiber,
54                current_suspend,
55                current_poll_cx,
56            }
57        };
58        future.await?;
59
60        Ok(slot.unwrap())
61    }
62
63    /// This is a helper function to call `resume` on the underlying
64    /// fiber while correctly managing thread-local data.
65    fn resume(&mut self, val: Result<(), ()>) -> Result<Result<(), ()>, ()> {
66        let async_cx = AsyncCx {
67            current_suspend: self.current_suspend,
68            current_poll_cx: self.current_poll_cx,
69        };
70        ASYNC_CX.set(&async_cx, || self.fiber.resume(val))
71    }
72}
73impl Future for FiberFuture<'_> {
74    type Output = Result<(), ()>;
75    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
76        unsafe {
77            let _reset = Reset(self.current_poll_cx, *self.current_poll_cx);
78            *self.current_poll_cx =
79                std::mem::transmute::<&mut Context<'_>, *mut Context<'static>>(cx);
80
81            match self.resume(Ok(())) {
82                Ok(ret) => Poll::Ready(ret),
83                Err(_) => Poll::Pending,
84            }
85        }
86    }
87}
88unsafe impl Send for FiberFuture<'_> {}
89unsafe impl Sync for FiberFuture<'_> {}
90
91type FiberSuspend = Suspend<Result<(), ()>, (), Result<(), ()>>;
92
93impl Drop for FiberFuture<'_> {
94    fn drop(&mut self) {
95        if !self.fiber.done() {
96            let result = self.resume(Err(()));
97            // This resumption with an error should always complete the
98            // fiber. While it's technically possible for host code to catch
99            // the trap and re-resume, we'd ideally like to signal that to
100            // callers that they shouldn't be doing that.
101            debug_assert!(result.is_ok());
102        }
103    }
104}
105
106/// Defines a TimeoutFiberFuture.
107///
108/// Not available on musl targets because it depends on signal-based timeout
109/// mechanisms (sigsetjmp/siglongjmp, timer_create) that are gated behind
110/// `not(target_env = "musl")` in the executor module.
111#[cfg(not(target_env = "musl"))]
112pub(crate) struct TimeoutFiberFuture<'a> {
113    fiber: Fiber<'a, Result<(), ()>, (), Result<(), ()>>,
114    current_suspend: *mut *const Suspend<Result<(), ()>, (), Result<(), ()>>,
115    current_poll_cx: *mut *mut Context<'static>,
116    deadline: std::time::SystemTime,
117}
118
119#[cfg(not(target_env = "musl"))]
120impl<'a> TimeoutFiberFuture<'a> {
121    /// Create a fiber to execute the given function.
122    ///
123    /// # Arguments
124    ///
125    /// * `func` - The function to execute.
126    ///
127    /// * `async_state` - Used to store asynchronous state at run time.
128    ///
129    /// * `deadline` - The deadline the function to be run.
130    ///
131    /// # Error
132    ///
133    /// If fail to create the fiber stack or the fiber fail to resume, then an error is returned.
134    pub(crate) async fn on_fiber<R>(
135        async_state: &AsyncState,
136        func: impl FnOnce() -> R + Send,
137        deadline: std::time::SystemTime,
138    ) -> Result<R, ()> {
139        let mut slot = None;
140
141        let future = {
142            let current_poll_cx = async_state.current_poll_cx.get();
143            let current_suspend = async_state.current_suspend.get();
144
145            let stack = FiberStack::new(2 << 20).map_err(|_e| ())?;
146            let slot = &mut slot;
147            let fiber = Fiber::new(stack, move |keep_going, suspend| {
148                keep_going?;
149
150                unsafe {
151                    let _reset = Reset(current_suspend, *current_suspend);
152                    *current_suspend = suspend;
153                    *slot = Some(func());
154                    Ok(())
155                }
156            })
157            .map_err(|_e| ())?;
158
159            TimeoutFiberFuture {
160                fiber,
161                current_suspend,
162                current_poll_cx,
163                deadline,
164            }
165        };
166
167        future.await?;
168
169        Ok(slot.unwrap())
170    }
171}
172
173#[cfg(not(target_env = "musl"))]
174impl Future for TimeoutFiberFuture<'_> {
175    type Output = Result<(), ()>;
176    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
177        unsafe {
178            crate::executor::init_signal_listen();
179
180            let _reset = Reset(self.current_poll_cx, *self.current_poll_cx);
181            *self.current_poll_cx =
182                std::mem::transmute::<&mut Context<'_>, *mut Context<'static>>(cx);
183            let async_cx = AsyncCx {
184                current_suspend: self.current_suspend,
185                current_poll_cx: self.current_poll_cx,
186            };
187
188            ASYNC_CX.set(&async_cx, || {
189                let timeout = match self.deadline.duration_since(std::time::SystemTime::now()) {
190                    Ok(timeout) => timeout.max(std::time::Duration::from_millis(100)),
191                    Err(_) => return Poll::Ready(Err(())),
192                };
193
194                let mut self_thread = libc::pthread_self();
195                let mut timerid: libc::timer_t = std::mem::zeroed();
196                let mut sev: libc::sigevent = std::mem::zeroed();
197                sev.sigev_notify = libc::SIGEV_SIGNAL;
198                sev.sigev_signo = crate::executor::timeout_signo();
199                sev.sigev_value.sival_ptr = &mut self_thread as *mut _ as *mut libc::c_void;
200
201                if libc::timer_create(libc::CLOCK_REALTIME, &mut sev, &mut timerid) < 0 {
202                    return Poll::Ready(Err(()));
203                }
204
205                let mut value: libc::itimerspec = std::mem::zeroed();
206                value.it_value.tv_sec = timeout.as_secs() as _;
207                value.it_value.tv_nsec = timeout.subsec_nanos() as _;
208                if libc::timer_settime(timerid, 0, &value, std::ptr::null_mut()) < 0 {
209                    libc::timer_delete(timerid);
210                    return Poll::Ready(Err(()));
211                }
212
213                let mut env: setjmp::sigjmp_buf = std::mem::zeroed();
214                let jmp_state = crate::executor::JmpState {
215                    sigjmp_buf: &mut env,
216                };
217
218                crate::executor::JMP_BUF.set(&jmp_state, || {
219                    if setjmp::sigsetjmp(&mut env, 1) == 0 {
220                        let r = match self.as_ref().fiber.resume(Ok(())) {
221                            Ok(ret) => Poll::Ready(ret),
222                            Err(_) => Poll::Pending,
223                        };
224                        libc::timer_delete(timerid);
225                        r
226                    } else {
227                        libc::timer_delete(timerid);
228                        Poll::Ready(Err(()))
229                    }
230                })
231            })
232        }
233    }
234}
235#[cfg(not(target_env = "musl"))]
236unsafe impl Send for TimeoutFiberFuture<'_> {}
237#[cfg(not(target_env = "musl"))]
238unsafe impl Sync for TimeoutFiberFuture<'_> {}
239
240#[cfg(not(target_env = "musl"))]
241impl Drop for TimeoutFiberFuture<'_> {
242    fn drop(&mut self) {
243        if !self.fiber.done() {
244            let async_cx = AsyncCx {
245                current_suspend: self.current_suspend,
246                current_poll_cx: self.current_poll_cx,
247            };
248            let result = ASYNC_CX.set(&async_cx, || self.fiber.resume(Err(())));
249            // This resumption with an error should always complete the
250            // fiber. While it's technically possible for host code to catch
251            // the trap and re-resume, we'd ideally like to signal that to
252            // callers that they shouldn't be doing that.
253            debug_assert!(result.is_ok());
254        }
255    }
256}
257
258scoped_tls::scoped_thread_local!(static ASYNC_CX: AsyncCx);
259
260/// Defines a async state that contains the pointer to current poll context and current suspend.
261#[derive(Debug)]
262pub struct AsyncState {
263    current_suspend: std::cell::UnsafeCell<*const FiberSuspend>,
264    current_poll_cx: std::cell::UnsafeCell<*mut Context<'static>>,
265}
266impl Default for AsyncState {
267    fn default() -> Self {
268        Self::new()
269    }
270}
271impl AsyncState {
272    /// Creates a new async state.
273    pub fn new() -> Self {
274        AsyncState {
275            current_suspend: std::cell::UnsafeCell::new(std::ptr::null()),
276            current_poll_cx: std::cell::UnsafeCell::new(std::ptr::null_mut()),
277        }
278    }
279
280    /// Returns an async execution context.
281    ///
282    /// If the pointer of poll context is null, then None is returned.
283    pub fn async_cx(&self) -> Option<AsyncCx> {
284        let poll_cx_box_ptr = self.current_poll_cx.get();
285        if poll_cx_box_ptr.is_null() {
286            return None;
287        }
288        let poll_cx_inner_ptr = unsafe { *poll_cx_box_ptr };
289        if poll_cx_inner_ptr.is_null() {
290            return None;
291        }
292
293        Some(AsyncCx {
294            current_suspend: self.current_suspend.get(),
295            current_poll_cx: poll_cx_box_ptr,
296        })
297    }
298}
299unsafe impl Send for AsyncState {}
300unsafe impl Sync for AsyncState {}
301
302/// Defines an async execution context.
303#[derive(Debug, Clone, Copy)]
304pub struct AsyncCx {
305    current_suspend: *mut *const Suspend<Result<(), ()>, (), Result<(), ()>>,
306    current_poll_cx: *mut *mut Context<'static>,
307}
308impl Default for AsyncCx {
309    fn default() -> Self {
310        Self::new()
311    }
312}
313impl AsyncCx {
314    /// Creates a new async execution context.
315    pub fn new() -> Self {
316        ASYNC_CX.with(|async_cx| *async_cx)
317    }
318
319    /// Runs a future to completion.
320    ///
321    /// # Arguments
322    ///
323    /// * `future` - The future to run.
324    ///
325    /// # Error
326    ///
327    /// If fail to run, then an error is returned.
328    pub(crate) unsafe fn block_on<U>(
329        &self,
330        mut future: Pin<&mut (dyn Future<Output = U> + Send)>,
331    ) -> Result<U, ()> {
332        let suspend = *self.current_suspend;
333        let _reset = Reset(self.current_suspend, suspend);
334        *self.current_suspend = ptr::null();
335        assert!(!suspend.is_null());
336
337        loop {
338            let future_result = {
339                let poll_cx = *self.current_poll_cx;
340                let _reset = Reset(self.current_poll_cx, poll_cx);
341                *self.current_poll_cx = ptr::null_mut();
342                assert!(!poll_cx.is_null());
343                future.as_mut().poll(&mut *poll_cx)
344            };
345
346            match future_result {
347                Poll::Ready(t) => break Ok(t),
348                Poll::Pending => {}
349            }
350            let res = (*suspend).suspend(());
351            res?;
352        }
353    }
354}
355
356struct Reset<T: Copy>(*mut T, T);
357impl<T: Copy> Drop for Reset<T> {
358    fn drop(&mut self) {
359        unsafe {
360            *self.0 = self.1;
361        }
362    }
363}