simple_async/
task.rs

1use std::{
2    cell::{RefCell, UnsafeCell},
3    future::Future,
4    pin::Pin,
5    rc::Rc,
6    sync::mpsc::Sender,
7    task::{Context, Poll},
8};
9
10use crate::waker::from_task;
11
12/// Holds state of a single top-level task that can be run by an executor
13pub struct Task {
14    /// The actual storage needed to execute the Task's code
15    future: UnsafeCell<Box<dyn Future<Output = ()>>>,
16    /// A sender that can be used to re-queue this Task to the executor
17    pub(crate) task_queue: Sender<Rc<Task>>,
18}
19
20impl Task {
21    /// Creates a new Task instance
22    pub(crate) fn new(
23        future: impl Future<Output = ()> + 'static,
24        task_queue: Sender<Rc<Self>>,
25    ) -> Self {
26        Self {
27            future: UnsafeCell::new(Box::new(future)),
28            task_queue,
29        }
30    }
31
32    /// Continues running the Task to the end or the next yielding point
33    pub(crate) fn poll(self: Rc<Self>) -> Poll<()> {
34        // Accessing the UnsafeCell is safe, a Task will only ever be accessed from a single point in code in a single thread.
35        let future = unsafe { &mut *self.future.get() }.as_mut();
36        // Pinning is safe, Tasks are never moved out of the Rc by the executor
37        let pin = unsafe { Pin::new_unchecked(future) };
38
39        let task_sender = self.task_queue.clone();
40
41        let waker = from_task(self);
42        let mut context = Context::from_waker(&waker);
43
44        CURRENT_TASK_SENDER.with(|cell| {
45            cell.replace(Some(task_sender));
46            let res = pin.poll(&mut context);
47            cell.replace(None);
48            res
49        })
50    }
51
52    /// Spawns a new [`Task`] that will run on the same executor as the current [`Task`]
53    ///
54    /// # Panics
55    /// This function panics if it is called from a context outside an executor or an async function.
56    pub fn spawn(future: impl Future<Output = ()> + 'static) {
57        let task_sender = CURRENT_TASK_SENDER.with(|cell| {
58            cell.borrow()
59                .as_ref()
60                .expect("Task::spawn() called from outside an executor")
61                .clone()
62        });
63
64        let task_sender2 = task_sender.clone();
65
66        let task = Self::new(future, task_sender);
67        task_sender2.send(Rc::new(task)).unwrap();
68    }
69}
70
71thread_local! {
72    /// This variable holds a [`Sender`] that can be used to enqueue a new [`Task`] into the [`Executor`](crate::executor::Executor) that is currently running
73    pub(crate) static CURRENT_TASK_SENDER: RefCell<Option<Sender<Rc<Task>>>> = RefCell::new(None);
74}
simple_async/task.rs

simple_async/
task.rs
