1use fiber_for_wasmedge::{Fiber, FiberStack, Suspend};
4use std::{
5 future::Future,
6 pin::Pin,
7 ptr,
8 task::{Context, Poll},
9};
10
11pub(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 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 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 debug_assert!(result.is_ok());
102 }
103 }
104}
105
106#[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 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 debug_assert!(result.is_ok());
254 }
255 }
256}
257
258scoped_tls::scoped_thread_local!(static ASYNC_CX: AsyncCx);
259
260#[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 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 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#[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 pub fn new() -> Self {
316 ASYNC_CX.with(|async_cx| *async_cx)
317 }
318
319 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}