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::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
126struct SharedTaskState {
127 sleep_state: AtomicU32,
129 inter_task_stream: inter_task_wakeup::WakerState,
130
131 waitables: TryLock<BTreeMap<u32, CabiWaitable>>,
144
145 waitable_set: TryLock<Option<WaitableSet>>,
150}
151
152struct CabiWaitable {
154 callback: unsafe extern "C" fn(*mut c_void, u32),
155 callback_ptr: *mut c_void,
156}
157
158unsafe impl Send for CabiWaitable {}
159
160impl TaskState<'_> {
161 fn new(future: BoxFuture<'_>) -> TaskState<'_> {
162 let shared = Arc::new(SharedTaskState {
163 sleep_state: AtomicU32::new(0),
164 inter_task_stream: Default::default(),
165 waitables: Default::default(),
166 waitable_set: Default::default(),
167 });
168 TaskState {
169 waker: shared.clone().into(),
170 shared,
171 tasks: spawn::Tasks::new(future),
172 inter_task_wakeup: Default::default(),
173 }
174 }
175
176 fn remaining_work(&self) -> bool {
177 !self.shared.waitables.try_lock().unwrap().is_empty()
178 }
179
180 fn callback(&mut self, event0: u32, event1: u32, event2: u32) -> CallbackCode {
183 match event0 {
184 EVENT_NONE => rtdebug!("EVENT_NONE"),
185 EVENT_SUBTASK => rtdebug!("EVENT_SUBTASK({event1:#x}, {event2:#x})"),
186 EVENT_STREAM_READ => rtdebug!("EVENT_STREAM_READ({event1:#x}, {event2:#x})"),
187 EVENT_STREAM_WRITE => rtdebug!("EVENT_STREAM_WRITE({event1:#x}, {event2:#x})"),
188 EVENT_FUTURE_READ => rtdebug!("EVENT_FUTURE_READ({event1:#x}, {event2:#x})"),
189 EVENT_FUTURE_WRITE => rtdebug!("EVENT_FUTURE_WRITE({event1:#x}, {event2:#x})"),
190 EVENT_CANCEL => {
191 rtdebug!("EVENT_CANCEL");
192
193 return CallbackCode::Exit;
198 }
199 _ => unreachable!(),
200 }
201
202 self.with_p3_task_set(|me| {
203 me.shared
206 .sleep_state
207 .store(SLEEP_STATE_WOKEN, Ordering::Relaxed);
208
209 if event0 != EVENT_NONE {
217 me.deliver_waitable_event(event1, event2)
218 }
219
220 me.cancel_inter_task_stream_read();
224
225 loop {
226 let mut context = Context::from_waker(&me.waker);
227
228 me.shared
233 .sleep_state
234 .store(SLEEP_STATE_POLLING, Ordering::Relaxed);
235
236 let poll = me.tasks.poll_next(&mut context);
238
239 match poll {
240 Poll::Ready(()) => {
245 assert!(me.tasks.is_empty());
246 if me.remaining_work() {
247 let set = me.shared.waitable_set.try_lock().unwrap();
248 let waitable = set.as_ref().unwrap().as_raw();
249 break CallbackCode::Wait(waitable);
250 } else {
251 break CallbackCode::Exit;
252 }
253 }
254
255 Poll::Pending => {
260 assert!(!me.tasks.is_empty());
261 if me.shared.sleep_state.load(Ordering::Relaxed) == SLEEP_STATE_WOKEN {
262 if me.remaining_work() {
263 let (event0, event1, event2) = {
264 let set = me.shared.waitable_set.try_lock().unwrap();
265 set.as_ref().unwrap().poll()
266 };
267 if event0 != EVENT_NONE {
268 me.deliver_waitable_event(event1, event2);
269 continue;
270 }
271 }
272 break CallbackCode::Yield;
273 }
274
275 me.shared
279 .sleep_state
280 .store(SLEEP_STATE_SLEEPING, Ordering::Relaxed);
281 me.read_inter_task_stream();
282 let set = me.shared.waitable_set.try_lock().unwrap();
283 let waitable = set.as_ref().unwrap().as_raw();
284 break CallbackCode::Wait(waitable);
285 }
286 }
287 }
288 })
289 }
290
291 fn deliver_waitable_event(&mut self, waitable: u32, code: u32) {
295 WaitableSet::remove_waitable_from_all_sets(waitable);
296
297 if self
298 .inter_task_wakeup
299 .consume_waitable_event(waitable, code)
300 {
301 return;
302 }
303
304 let c = {
305 let mut waitables = self.shared.waitables.try_lock().unwrap();
306 waitables.remove(&waitable).unwrap()
307 };
308 unsafe {
309 (c.callback)(c.callback_ptr, code);
310 }
311 }
312
313 fn with_p3_task_set<R>(&mut self, f: impl FnOnce(&mut Self) -> R) -> R {
314 struct ResetTask(*mut cabi::wasip3_task);
318 impl Drop for ResetTask {
319 fn drop(&mut self) {
320 unsafe {
321 cabi::wasip3_task_set(self.0);
322 }
323 }
324 }
325 let shared_raw: *const SharedTaskState = &*self.shared;
328 let mut wasip3_task = cabi::wasip3_task_v2 {
329 v1: cabi::wasip3_task {
330 ptr: shared_raw.cast_mut().cast(),
331 version: cabi::WASIP3_TASK_V2,
332 waitable_register: SharedTaskState::CABI_VTABLE.waitable_register,
333 waitable_unregister: SharedTaskState::CABI_VTABLE.waitable_unregister,
334 },
335 vtable: &SharedTaskState::CABI_VTABLE,
336 };
337
338 let wasip3_task: *mut cabi::wasip3_task_v2 = &mut wasip3_task;
343 let prev = unsafe { cabi::wasip3_task_set(wasip3_task.cast::<cabi::wasip3_task>()) };
344 let _reset = ResetTask(prev);
345
346 f(self)
347 }
348}
349
350impl Drop for TaskState<'_> {
351 fn drop(&mut self) {
352 self.cancel_inter_task_stream_read();
355
356 if !self.tasks.is_empty() {
362 self.with_p3_task_set(|me| {
363 me.tasks = Default::default();
364 })
365 }
366 }
367}
368
369impl SharedTaskState {
370 const CABI_VTABLE: cabi::wasip3_task_vtable = cabi::wasip3_task_vtable {
371 waitable_register: Self::cabi_waitable_register,
372 waitable_unregister: Self::cabi_waitable_unregister,
373 drop: Self::cabi_drop,
374 clone: Self::cabi_clone,
375 };
376
377 fn add_waitable(&self, waitable: u32) {
379 let mut set = self.waitable_set.try_lock().unwrap();
380 set.get_or_insert_with(WaitableSet::new).join(waitable);
381 }
382
383 fn waitable_register(
385 &self,
386 waitable: u32,
387 callback: unsafe extern "C" fn(*mut c_void, u32),
388 callback_ptr: *mut c_void,
389 ) -> *mut c_void {
390 self.add_waitable(waitable);
391 let mut waitables = self.waitables.try_lock().unwrap();
392 let c = CabiWaitable {
393 callback,
394 callback_ptr,
395 };
396 match waitables.insert(waitable, c) {
397 Some(prev) => prev.callback_ptr,
398 None => ptr::null_mut(),
399 }
400 }
401
402 fn waitable_unregister(&self, waitable: u32) -> *mut c_void {
404 WaitableSet::remove_waitable_from_all_sets(waitable);
405 let mut waitables = self.waitables.try_lock().unwrap();
406 match waitables.remove(&waitable) {
407 Some(prev) => prev.callback_ptr,
408 None => ptr::null_mut(),
409 }
410 }
411
412 unsafe fn cabi_to_self(ptr: *mut c_void) -> ManuallyDrop<Arc<SharedTaskState>> {
415 unsafe { ManuallyDrop::new(Arc::from_raw(ptr.cast::<SharedTaskState>())) }
416 }
417
418 unsafe extern "C" fn cabi_waitable_register(
419 ptr: *mut c_void,
420 waitable: u32,
421 callback: unsafe extern "C" fn(*mut c_void, u32),
422 callback_ptr: *mut c_void,
423 ) -> *mut c_void {
424 let me = unsafe { Self::cabi_to_self(ptr) };
425 me.waitable_register(waitable, callback, callback_ptr)
426 }
427
428 unsafe extern "C" fn cabi_waitable_unregister(ptr: *mut c_void, waitable: u32) -> *mut c_void {
429 let me = unsafe { Self::cabi_to_self(ptr) };
430 me.waitable_unregister(waitable)
431 }
432
433 unsafe extern "C" fn cabi_clone(ptr: *mut c_void) -> *mut c_void {
434 let me = unsafe { Self::cabi_to_self(ptr) };
435 Arc::into_raw(Arc::clone(&me)).cast_mut().cast()
436 }
437
438 unsafe extern "C" fn cabi_drop(ptr: *mut c_void) {
439 let mut me = unsafe { Self::cabi_to_self(ptr) };
440 unsafe { ManuallyDrop::drop(&mut me) }
441 }
442}
443
444const SLEEP_STATE_POLLING: u32 = 0;
446const SLEEP_STATE_WOKEN: u32 = 1;
448const SLEEP_STATE_SLEEPING: u32 = 2;
452
453impl Wake for SharedTaskState {
454 fn wake(self: Arc<Self>) {
455 Self::wake_by_ref(&self)
456 }
457
458 fn wake_by_ref(self: &Arc<Self>) {
459 match self.sleep_state.swap(SLEEP_STATE_WOKEN, Ordering::Relaxed) {
460 SLEEP_STATE_POLLING | SLEEP_STATE_WOKEN => {}
463
464 other => {
467 assert_eq!(other, SLEEP_STATE_SLEEPING);
468 self.inter_task_stream.wake();
469 }
470 }
471 }
472}
473
474const EVENT_NONE: u32 = 0;
475const EVENT_SUBTASK: u32 = 1;
476const EVENT_STREAM_READ: u32 = 2;
477const EVENT_STREAM_WRITE: u32 = 3;
478const EVENT_FUTURE_READ: u32 = 4;
479const EVENT_FUTURE_WRITE: u32 = 5;
480const EVENT_CANCEL: u32 = 6;
481
482#[derive(PartialEq, Debug)]
483enum CallbackCode {
484 Exit,
485 Yield,
486 Wait(u32),
487}
488
489impl CallbackCode {
490 fn encode(self) -> u32 {
491 match self {
492 CallbackCode::Exit => 0,
493 CallbackCode::Yield => 1,
494 CallbackCode::Wait(waitable) => 2 | (waitable << 4),
495 }
496 }
497}
498
499const STATUS_STARTING: u32 = 0;
500const STATUS_STARTED: u32 = 1;
501const STATUS_RETURNED: u32 = 2;
502const STATUS_STARTED_CANCELLED: u32 = 3;
503const STATUS_RETURNED_CANCELLED: u32 = 4;
504
505const BLOCKED: u32 = 0xffff_ffff;
506const COMPLETED: u32 = 0x0;
507const DROPPED: u32 = 0x1;
508const CANCELLED: u32 = 0x2;
509
510#[derive(PartialEq, Debug, Copy, Clone)]
512enum ReturnCode {
513 Blocked,
515 Completed(u32),
517 Dropped(u32),
520 Cancelled(u32),
523}
524
525impl ReturnCode {
526 fn decode(val: u32) -> ReturnCode {
527 if val == BLOCKED {
528 return ReturnCode::Blocked;
529 }
530 let amt = val >> 4;
531 match val & 0xf {
532 COMPLETED => ReturnCode::Completed(amt),
533 DROPPED => ReturnCode::Dropped(amt),
534 CANCELLED => ReturnCode::Cancelled(amt),
535 _ => panic!("unknown return code {val:#x}"),
536 }
537 }
538}
539
540#[doc(hidden)]
548pub fn start_task(task: impl Future<Output = ()> + 'static) -> i32 {
549 let state = Box::into_raw(Box::new(TaskState::new(Box::pin(task))));
552
553 unsafe {
560 assert!(context_get().is_null());
561 context_set(state.cast());
562 callback(EVENT_NONE, 0, 0) as i32
563 }
564}
565
566#[doc(hidden)]
573pub unsafe fn callback(event0: u32, event1: u32, event2: u32) -> u32 {
574 let state = context_get().cast::<TaskState<'static>>();
578 assert!(!state.is_null());
579 unsafe {
580 context_set(ptr::null_mut());
581 }
582
583 unsafe {
588 let rc = (*state).callback(event0, event1, event2);
589 if rc == CallbackCode::Exit {
590 drop(Box::from_raw(state));
591 } else {
592 context_set(state.cast());
593 }
594 rtdebug!(" => (cb) {rc:?}");
595 rc.encode()
596 }
597}
598
599pub fn block_on<T: 'static>(future: impl Future<Output = T>) -> T {
605 let mut result = None;
606 let mut state = TaskState::new(Box::pin(async {
607 result = Some(future.await);
608 }));
609 let mut event = (EVENT_NONE, 0, 0);
610 loop {
611 match state.callback(event.0, event.1, event.2) {
612 CallbackCode::Exit => {
613 drop(state);
614 break result.unwrap();
615 }
616 CallbackCode::Yield => {
617 let set = state.shared.waitable_set.try_lock().unwrap();
618 event = set.as_ref().unwrap().poll()
619 }
620 CallbackCode::Wait(_) => {
621 let set = state.shared.waitable_set.try_lock().unwrap();
622 event = set.as_ref().unwrap().wait()
623 }
624 }
625 }
626}
627
628pub fn yield_blocking() -> bool {
649 extern_wasm! {
650 #[link(wasm_import_module = "$root")]
651 unsafe extern "C" {
652 #[link_name = "[thread-yield]"]
653 fn yield_() -> bool;
654 }
655 }
656
657 unsafe { !yield_() }
661}
662
663pub async fn yield_async() {
678 #[derive(Default)]
679 struct Yield {
680 yielded: bool,
681 }
682
683 impl Future for Yield {
684 type Output = ();
685
686 fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<()> {
687 if self.yielded {
688 Poll::Ready(())
689 } else {
690 self.yielded = true;
691 context.waker().wake_by_ref();
692 Poll::Pending
693 }
694 }
695 }
696
697 Yield::default().await;
698}
699
700pub fn backpressure_inc() {
702 extern_wasm! {
703 #[link(wasm_import_module = "$root")]
704 unsafe extern "C" {
705 #[link_name = "[backpressure-inc]"]
706 fn backpressure_inc();
707 }
708 }
709
710 unsafe { backpressure_inc() }
711}
712
713pub fn backpressure_dec() {
715 extern_wasm! {
716 #[link(wasm_import_module = "$root")]
717 unsafe extern "C" {
718 #[link_name = "[backpressure-dec]"]
719 fn backpressure_dec();
720 }
721 }
722
723 unsafe { backpressure_dec() }
724}
725
726fn context_get() -> *mut u8 {
727 extern_wasm! {
728 #[link(wasm_import_module = "$root")]
729 unsafe extern "C" {
730 #[link_name = "[context-get-0]"]
731 fn get() -> *mut u8;
732 }
733 }
734
735 unsafe { get() }
736}
737
738unsafe fn context_set(value: *mut u8) {
739 extern_wasm! {
740 #[link(wasm_import_module = "$root")]
741 unsafe extern "C" {
742 #[link_name = "[context-set-0]"]
743 fn set(value: *mut u8);
744 }
745 }
746
747 unsafe { set(value) }
748}
749
750#[doc(hidden)]
751pub struct TaskCancelOnDrop {
752 _priv: (),
753}
754
755impl TaskCancelOnDrop {
756 #[doc(hidden)]
757 pub fn new() -> TaskCancelOnDrop {
758 TaskCancelOnDrop { _priv: () }
759 }
760
761 #[doc(hidden)]
762 pub fn forget(self) {
763 mem::forget(self);
764 }
765}
766
767impl Drop for TaskCancelOnDrop {
768 fn drop(&mut self) {
769 extern_wasm! {
770 #[link(wasm_import_module = "[export]$root")]
771 unsafe extern "C" {
772 #[link_name = "[task-cancel]"]
773 fn cancel();
774 }
775 }
776
777 unsafe { cancel() }
778 }
779}