1#![deny(missing_docs)]
2
3use self::try_lock::TryLock;
4use alloc::boxed::Box;
5use alloc::collections::BTreeMap;
6use alloc::sync::Arc;
7use alloc::task::Wake;
8use core::ffi::c_void;
9use core::future::Future;
10use core::mem::{self, ManuallyDrop};
11use core::pin::Pin;
12use core::ptr;
13use core::sync::atomic::{AtomicU32, Ordering};
14use core::task::{Context, Poll, Waker};
15
16macro_rules! rtdebug {
17 ($($f:tt)*) => {
18 if false {
22 #[cfg(feature = "std")]
23 std::eprintln!($($f)*);
24 }
25 }
26}
27
28macro_rules! extern_wasm {
34 (
35 $(#[$extern_attr:meta])*
36 unsafe extern "C" {
37 $(
38 $(#[$func_attr:meta])*
39 $vis:vis fn $func_name:ident ( $($args:tt)* ) $(-> $ret:ty)?;
40 )*
41 }
42 ) => {
43 $(
44 #[cfg(not(target_family = "wasm"))]
45 #[allow(unused, reason = "dummy shim for non-wasm compilation, never invoked")]
46 $vis unsafe fn $func_name($($args)*) $(-> $ret)? {
47 unreachable!();
48 }
49 )*
50
51 #[cfg(target_family = "wasm")]
52 $(#[$extern_attr])*
53 unsafe extern "C" {
54 $(
55 $(#[$func_attr])*
56 $vis fn $func_name($($args)*) $(-> $ret)?;
57 )*
58 }
59 };
60}
61
62mod abi_buffer;
63mod cabi;
64mod error_context;
65mod future_support;
66#[cfg(feature = "futures-stream")]
67mod futures_stream;
68#[cfg(feature = "inter-task-wakeup")]
69mod inter_task_wakeup;
70mod stream_support;
71mod subtask;
72mod try_lock;
73#[cfg(feature = "inter-task-wakeup")]
74mod unit_stream;
75mod waitable;
76mod waitable_set;
77
78#[cfg(not(feature = "inter-task-wakeup"))]
79use inter_task_wakeup_disabled as inter_task_wakeup;
80#[cfg(not(feature = "inter-task-wakeup"))]
81mod inter_task_wakeup_disabled;
82
83use self::waitable_set::WaitableSet;
84pub use abi_buffer::*;
85pub use error_context::*;
86pub use future_support::*;
87#[cfg(feature = "futures-stream")]
88pub use futures_stream::*;
89pub use stream_support::*;
90#[doc(hidden)]
91pub use subtask::Subtask;
92#[cfg(feature = "inter-task-wakeup")]
93pub use unit_stream::*;
94
95type BoxFuture<'a> = Pin<Box<dyn Future<Output = ()> + 'a>>;
96
97#[cfg(feature = "async-spawn")]
98mod spawn;
99#[cfg(feature = "async-spawn")]
100pub use spawn::{JoinHandle, spawn_local};
101#[cfg(not(feature = "async-spawn"))]
102mod spawn_disabled;
103#[cfg(not(feature = "async-spawn"))]
104use spawn_disabled as spawn;
105
106struct TaskState<'a> {
109 tasks: spawn::Tasks<'a>,
114
115 shared: Arc<SharedTaskState>,
118
119 waker: Waker,
121
122 inter_task_wakeup: inter_task_wakeup::State,
124}
125
126#[derive(Default)]
127struct SharedTaskState {
128 sleep_state: AtomicU32,
130 inter_task_stream: inter_task_wakeup::WakerState,
131
132 waitables: TryLock<BTreeMap<u32, CabiWaitable>>,
145
146 waitable_set: TryLock<Option<WaitableSet>>,
151}
152
153struct CabiWaitable {
155 callback: unsafe extern "C" fn(*mut c_void, u32),
156 callback_ptr: *mut c_void,
157}
158
159unsafe impl Send for CabiWaitable {}
160
161impl<'a> TaskState<'a> {
162 fn new(future: impl Future<Output = ()> + 'a) -> TaskState<'a> {
163 let shared = Arc::new(SharedTaskState::default());
164 TaskState {
165 waker: shared.clone().into(),
166 shared,
167 tasks: spawn::Tasks::new(future),
168 inter_task_wakeup: Default::default(),
169 }
170 }
171
172 fn remaining_work(&self) -> bool {
173 !self.shared.waitables.try_lock().unwrap().is_empty()
174 }
175
176 fn callback(&mut self, event0: u32, event1: u32, event2: u32) -> CallbackCode {
179 match event0 {
180 EVENT_NONE => rtdebug!("EVENT_NONE"),
181 EVENT_SUBTASK => rtdebug!("EVENT_SUBTASK({event1:#x}, {event2:#x})"),
182 EVENT_STREAM_READ => rtdebug!("EVENT_STREAM_READ({event1:#x}, {event2:#x})"),
183 EVENT_STREAM_WRITE => rtdebug!("EVENT_STREAM_WRITE({event1:#x}, {event2:#x})"),
184 EVENT_FUTURE_READ => rtdebug!("EVENT_FUTURE_READ({event1:#x}, {event2:#x})"),
185 EVENT_FUTURE_WRITE => rtdebug!("EVENT_FUTURE_WRITE({event1:#x}, {event2:#x})"),
186
187 EVENT_CANCEL => {
192 rtdebug!("EVENT_CANCEL");
193 self.with_p3_task_set(|me| me.tasks.cancel_root());
194 }
195 _ => unreachable!(),
196 }
197
198 self.with_p3_task_set(|me| {
199 me.shared
202 .sleep_state
203 .store(SLEEP_STATE_WOKEN, Ordering::Relaxed);
204
205 if event0 != EVENT_NONE && event0 != EVENT_CANCEL {
213 me.deliver_waitable_event(event1, event2)
214 }
215
216 me.cancel_inter_task_stream_read();
220
221 loop {
222 let mut context = Context::from_waker(&me.waker);
223
224 me.shared
229 .sleep_state
230 .store(SLEEP_STATE_POLLING, Ordering::Relaxed);
231
232 let poll = me.tasks.poll_next(&mut context);
234
235 match poll {
236 Poll::Ready(()) => {
241 assert!(me.tasks.is_empty());
242 if me.remaining_work() {
243 let set = me.shared.waitable_set.try_lock().unwrap();
244 let waitable = set.as_ref().unwrap().as_raw();
245 break CallbackCode::Wait(waitable);
246 } else {
247 break CallbackCode::Exit;
248 }
249 }
250
251 Poll::Pending => {
256 assert!(!me.tasks.is_empty());
257 if me.shared.sleep_state.load(Ordering::Relaxed) == SLEEP_STATE_WOKEN {
258 if me.remaining_work() {
259 let (event0, event1, event2) = {
260 let set = me.shared.waitable_set.try_lock().unwrap();
261 set.as_ref().unwrap().poll()
262 };
263 if event0 != EVENT_NONE {
264 me.deliver_waitable_event(event1, event2);
265 continue;
266 }
267 }
268 break CallbackCode::Yield;
269 }
270
271 me.shared
275 .sleep_state
276 .store(SLEEP_STATE_SLEEPING, Ordering::Relaxed);
277 me.read_inter_task_stream();
278 let set = me.shared.waitable_set.try_lock().unwrap();
279 let waitable = set.as_ref().unwrap().as_raw();
280 break CallbackCode::Wait(waitable);
281 }
282 }
283 }
284 })
285 }
286
287 fn deliver_waitable_event(&mut self, waitable: u32, code: u32) {
291 WaitableSet::remove_waitable_from_all_sets(waitable);
292
293 if self
294 .inter_task_wakeup
295 .consume_waitable_event(waitable, code)
296 {
297 return;
298 }
299
300 let c = {
301 let mut waitables = self.shared.waitables.try_lock().unwrap();
302 waitables.remove(&waitable).unwrap()
303 };
304 unsafe {
305 (c.callback)(c.callback_ptr, code);
306 }
307 }
308
309 fn with_p3_task_set<R>(&mut self, f: impl FnOnce(&mut Self) -> R) -> R {
310 struct ResetTask(*mut cabi::wasip3_task);
314 impl Drop for ResetTask {
315 fn drop(&mut self) {
316 unsafe {
317 cabi::wasip3_task_set(self.0);
318 }
319 }
320 }
321 let shared_raw: *const SharedTaskState = &*self.shared;
324 let mut wasip3_task = cabi::wasip3_task_v3 {
325 v1: cabi::wasip3_task {
326 ptr: shared_raw.cast_mut().cast(),
327 version: cabi::WASIP3_TASK_V3,
328 waitable_register: SharedTaskState::CABI_VTABLE.v2.waitable_register,
329 waitable_unregister: SharedTaskState::CABI_VTABLE.v2.waitable_unregister,
330 },
331 vtable: &SharedTaskState::CABI_VTABLE,
332 };
333
334 let wasip3_task: *mut cabi::wasip3_task_v3 = &mut wasip3_task;
339 let prev = unsafe { cabi::wasip3_task_set(wasip3_task.cast::<cabi::wasip3_task>()) };
340 let _reset = ResetTask(prev);
341
342 f(self)
343 }
344}
345
346impl Drop for TaskState<'_> {
347 fn drop(&mut self) {
348 self.shared
352 .sleep_state
353 .store(SLEEP_STATE_WOKEN, Ordering::Relaxed);
354
355 self.cancel_inter_task_stream_read();
358
359 if !self.tasks.is_empty() {
365 self.with_p3_task_set(|me| {
366 me.tasks = Default::default();
367 })
368 }
369 }
370}
371
372impl SharedTaskState {
373 const CABI_VTABLE: cabi::wasip3_task_vtable_v3 = cabi::wasip3_task_vtable_v3 {
374 v2: cabi::wasip3_task_vtable {
375 waitable_register: Self::cabi_waitable_register,
376 waitable_unregister: Self::cabi_waitable_unregister,
377 drop: Self::cabi_drop,
378 clone: Self::cabi_clone,
379 },
380 #[cfg(feature = "async-spawn")]
381 rust_spawn: Some(Self::cabi_rust_spawn),
382 #[cfg(not(feature = "async-spawn"))]
383 rust_spawn: None,
384 };
385
386 fn add_waitable(&self, waitable: u32) {
388 let mut set = self.waitable_set.try_lock().unwrap();
389 set.get_or_insert_with(WaitableSet::new).join(waitable);
390 }
391
392 fn waitable_register(
394 &self,
395 waitable: u32,
396 callback: unsafe extern "C" fn(*mut c_void, u32),
397 callback_ptr: *mut c_void,
398 ) -> *mut c_void {
399 self.add_waitable(waitable);
400 let mut waitables = self.waitables.try_lock().unwrap();
401 let c = CabiWaitable {
402 callback,
403 callback_ptr,
404 };
405 match waitables.insert(waitable, c) {
406 Some(prev) => prev.callback_ptr,
407 None => ptr::null_mut(),
408 }
409 }
410
411 fn waitable_unregister(&self, waitable: u32) -> *mut c_void {
413 WaitableSet::remove_waitable_from_all_sets(waitable);
414 let mut waitables = self.waitables.try_lock().unwrap();
415 match waitables.remove(&waitable) {
416 Some(prev) => prev.callback_ptr,
417 None => ptr::null_mut(),
418 }
419 }
420
421 unsafe fn cabi_to_self(ptr: *mut c_void) -> ManuallyDrop<Arc<SharedTaskState>> {
424 unsafe { ManuallyDrop::new(Arc::from_raw(ptr.cast::<SharedTaskState>())) }
425 }
426
427 unsafe extern "C" fn cabi_waitable_register(
428 ptr: *mut c_void,
429 waitable: u32,
430 callback: unsafe extern "C" fn(*mut c_void, u32),
431 callback_ptr: *mut c_void,
432 ) -> *mut c_void {
433 let me = unsafe { Self::cabi_to_self(ptr) };
434 me.waitable_register(waitable, callback, callback_ptr)
435 }
436
437 unsafe extern "C" fn cabi_waitable_unregister(ptr: *mut c_void, waitable: u32) -> *mut c_void {
438 let me = unsafe { Self::cabi_to_self(ptr) };
439 me.waitable_unregister(waitable)
440 }
441
442 unsafe extern "C" fn cabi_clone(ptr: *mut c_void) -> *mut c_void {
443 let me = unsafe { Self::cabi_to_self(ptr) };
444 Arc::into_raw(Arc::clone(&me)).cast_mut().cast()
445 }
446
447 unsafe extern "C" fn cabi_drop(ptr: *mut c_void) {
448 let mut me = unsafe { Self::cabi_to_self(ptr) };
449 unsafe { ManuallyDrop::drop(&mut me) }
450 }
451
452 #[cfg(feature = "async-spawn")]
453 unsafe fn cabi_rust_spawn(_ptr: *mut c_void, task: Box<dyn Future<Output = ()>>) {
454 spawn::push(Pin::from(task))
455 }
456}
457
458const SLEEP_STATE_POLLING: u32 = 0;
460const SLEEP_STATE_WOKEN: u32 = 1;
463const SLEEP_STATE_SLEEPING: u32 = 2;
467
468impl Wake for SharedTaskState {
469 fn wake(self: Arc<Self>) {
470 Self::wake_by_ref(&self)
471 }
472
473 fn wake_by_ref(self: &Arc<Self>) {
474 match self.sleep_state.swap(SLEEP_STATE_WOKEN, Ordering::Relaxed) {
475 SLEEP_STATE_POLLING | SLEEP_STATE_WOKEN => {}
478
479 other => {
482 assert_eq!(other, SLEEP_STATE_SLEEPING);
483 self.inter_task_stream.wake();
484 }
485 }
486 }
487}
488
489const EVENT_NONE: u32 = 0;
490const EVENT_SUBTASK: u32 = 1;
491const EVENT_STREAM_READ: u32 = 2;
492const EVENT_STREAM_WRITE: u32 = 3;
493const EVENT_FUTURE_READ: u32 = 4;
494const EVENT_FUTURE_WRITE: u32 = 5;
495const EVENT_CANCEL: u32 = 6;
496
497#[derive(PartialEq, Debug)]
498enum CallbackCode {
499 Exit,
500 Yield,
501 Wait(u32),
502}
503
504impl CallbackCode {
505 fn encode(self) -> u32 {
506 match self {
507 CallbackCode::Exit => 0,
508 CallbackCode::Yield => 1,
509 CallbackCode::Wait(waitable) => 2 | (waitable << 4),
510 }
511 }
512}
513
514const STATUS_STARTING: u32 = 0;
515const STATUS_STARTED: u32 = 1;
516const STATUS_RETURNED: u32 = 2;
517const STATUS_STARTED_CANCELLED: u32 = 3;
518const STATUS_RETURNED_CANCELLED: u32 = 4;
519
520const BLOCKED: u32 = 0xffff_ffff;
521const COMPLETED: u32 = 0x0;
522const DROPPED: u32 = 0x1;
523const CANCELLED: u32 = 0x2;
524
525#[derive(PartialEq, Debug, Copy, Clone)]
527enum ReturnCode {
528 Blocked,
530 Completed(u32),
532 Dropped(u32),
535 Cancelled(u32),
538}
539
540impl ReturnCode {
541 fn decode(val: u32) -> ReturnCode {
542 if val == BLOCKED {
543 return ReturnCode::Blocked;
544 }
545 let amt = val >> 4;
546 match val & 0xf {
547 COMPLETED => ReturnCode::Completed(amt),
548 DROPPED => ReturnCode::Dropped(amt),
549 CANCELLED => ReturnCode::Cancelled(amt),
550 _ => panic!("unknown return code {val:#x}"),
551 }
552 }
553}
554
555#[cfg(not(target_env = "p3"))]
556mod task_state {
557 pub fn get() -> *mut u8 {
558 extern_wasm! {
559 #[link(wasm_import_module = "$root")]
560 unsafe extern "C" {
561 #[link_name = "[context-get-0]"]
562 fn get() -> *mut u8;
563 }
564 }
565
566 unsafe { get() }
567 }
568
569 pub unsafe fn set(value: *mut u8) {
570 extern_wasm! {
571 #[link(wasm_import_module = "$root")]
572 unsafe extern "C" {
573 #[link_name = "[context-set-0]"]
574 fn set(value: *mut u8);
575 }
576 }
577
578 unsafe { set(value) }
579 }
580}
581
582#[cfg(all(target_env = "p3", feature = "std"))]
583mod task_state {
584 std::thread_local!(static TASK_STATE: std::cell::Cell<*mut u8> = std::cell::Cell::new(std::ptr::null_mut()));
585
586 pub fn get() -> *mut u8 {
587 TASK_STATE.with(|c| c.get())
588 }
589
590 pub unsafe fn set(value: *mut u8) {
591 TASK_STATE.with(|c| c.set(value))
592 }
593}
594
595#[cfg(all(target_env = "p3", not(feature = "std")))]
596mod wasip3_context;
597#[cfg(all(target_env = "p3", not(feature = "std")))]
598mod task_state {
599
600 pub fn get() -> *mut u8 {
601 unsafe { super::wasip3_context::get() }
602 }
603
604 pub unsafe fn set(value: *mut u8) {
605 unsafe { super::wasip3_context::set(value) }
606 }
607}
608
609#[doc(hidden)]
617pub fn start_task(task: impl Future<Output = ()> + 'static) -> i32 {
618 let state = Box::into_raw(Box::new(TaskState::new(task)));
621
622 unsafe {
629 assert!(task_state::get().is_null());
630 task_state::set(state.cast());
631 callback(EVENT_NONE, 0, 0) as i32
632 }
633}
634
635#[doc(hidden)]
642pub unsafe fn callback(event0: u32, event1: u32, event2: u32) -> u32 {
643 let state = task_state::get().cast::<TaskState<'static>>();
647 assert!(!state.is_null());
648 unsafe {
649 task_state::set(ptr::null_mut());
650 }
651
652 unsafe {
657 let rc = (*state).callback(event0, event1, event2);
658 if rc == CallbackCode::Exit {
659 drop(Box::from_raw(state));
660 } else {
661 task_state::set(state.cast());
662 }
663 rtdebug!(" => (cb) {rc:?}");
664 rc.encode()
665 }
666}
667
668pub fn block_on<T>(future: impl Future<Output = T>) -> T {
673 let mut result = None;
674 let mut state = TaskState::new(async {
675 result = Some(future.await);
676 });
677 let mut event = (EVENT_NONE, 0, 0);
678 loop {
679 match state.callback(event.0, event.1, event.2) {
680 CallbackCode::Exit => {
681 drop(state);
682 break result.unwrap();
683 }
684 CallbackCode::Yield => {
689 event = (EVENT_NONE, 0, 0);
690 }
691 CallbackCode::Wait(_) => {
692 let set = state.shared.waitable_set.try_lock().unwrap();
693 event = set.as_ref().unwrap().wait()
694 }
695 }
696 }
697}
698
699pub fn yield_blocking() -> bool {
720 extern_wasm! {
721 #[link(wasm_import_module = "$root")]
722 unsafe extern "C" {
723 #[link_name = "[thread-yield]"]
724 fn yield_() -> bool;
725 }
726 }
727
728 unsafe { !yield_() }
732}
733
734pub async fn yield_async() {
749 #[derive(Default)]
750 struct Yield {
751 yielded: bool,
752 }
753
754 impl Future for Yield {
755 type Output = ();
756
757 fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<()> {
758 if self.yielded {
759 Poll::Ready(())
760 } else {
761 self.yielded = true;
762 context.waker().wake_by_ref();
763 Poll::Pending
764 }
765 }
766 }
767
768 Yield::default().await;
769}
770
771pub fn backpressure_inc() {
773 extern_wasm! {
774 #[link(wasm_import_module = "$root")]
775 unsafe extern "C" {
776 #[link_name = "[backpressure-inc]"]
777 fn backpressure_inc();
778 }
779 }
780
781 unsafe { backpressure_inc() }
782}
783
784pub fn backpressure_dec() {
786 extern_wasm! {
787 #[link(wasm_import_module = "$root")]
788 unsafe extern "C" {
789 #[link_name = "[backpressure-dec]"]
790 fn backpressure_dec();
791 }
792 }
793
794 unsafe { backpressure_dec() }
795}
796
797#[doc(hidden)]
798pub struct TaskCancelOnDrop {
799 _priv: (),
800}
801
802impl TaskCancelOnDrop {
803 #[doc(hidden)]
804 pub fn new() -> TaskCancelOnDrop {
805 TaskCancelOnDrop { _priv: () }
806 }
807
808 #[doc(hidden)]
809 pub fn forget(self) {
810 mem::forget(self);
811 }
812}
813
814impl Drop for TaskCancelOnDrop {
815 fn drop(&mut self) {
816 extern_wasm! {
817 #[link(wasm_import_module = "[export]$root")]
818 unsafe extern "C" {
819 #[link_name = "[task-cancel]"]
820 fn cancel();
821 }
822 }
823
824 unsafe { cancel() }
825 }
826}