ndless_async/
yield_now.rs

1use alloc::rc::Rc;
2use core::cell::{Cell, RefCell};
3use core::future::Future;
4use core::pin::Pin;
5
6use futures_util::task::{AtomicWaker, Context, Poll};
7
8use ndless::prelude::*;
9
10struct WakerData {
11	waker: AtomicWaker,
12	done: Cell<bool>,
13}
14
15#[derive(Default)]
16pub(crate) struct YieldListener {
17	wakers: RefCell<Vec<Rc<WakerData>>>,
18}
19
20impl YieldListener {
21	pub(crate) fn poll(&self) {
22		let mut wakers = self.wakers.borrow_mut();
23		wakers.retain(|waker| Rc::strong_count(waker) > 1);
24		wakers.iter_mut().for_each(|waker| {
25			waker.done.set(true);
26			waker.waker.wake();
27		})
28	}
29	pub(crate) fn yield_now(&self) -> Yield {
30		let waker = Rc::new(WakerData {
31			done: Cell::new(false),
32			waker: AtomicWaker::new(),
33		});
34		let mut wakers = self.wakers.borrow_mut();
35		wakers.push(waker.clone());
36		Yield { waker }
37	}
38}
39
40/// Allows other tasks to run. See
41/// [`AsyncListeners::yield_now`][crate::task::AsyncListeners::yield_now]
42/// for more details.
43pub struct Yield {
44	waker: Rc<WakerData>,
45}
46
47impl Future for Yield {
48	type Output = ();
49
50	fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
51		if self.waker.done.get() {
52			Poll::Ready(())
53		} else {
54			self.waker.waker.register(cx.waker());
55			Poll::Pending
56		}
57	}
58}