use crate::{Join, JoinHandle, SchedInfo};
use alloc::boxed::Box;
use alloc::rc::{Rc, Weak};
use core::cell::RefCell;
use core::future::Future;
use core::pin::Pin;
use core::task::{Context, Poll};
#[derive(Copy, Clone, Debug)]
pub enum TaskStat {
Init,
Running,
Aborted,
End,
}
pub(crate) trait Task {
fn get_id(&self) -> u64;
fn abort(&mut self);
fn run(&mut self, ctx: &mut Context<'_>) -> Poll<()>;
}
pub(crate) struct TaskImpl<T: Future> {
id: u64,
stat: TaskStat,
future: T,
handle: Weak<dyn Join<T::Output>>,
}
impl<T: Future + 'static> TaskImpl<T> {
pub(crate) fn with_info(
future: T,
info: &Rc<RefCell<SchedInfo>>,
) -> (Box<dyn Task>, JoinHandle<T::Output>) {
let id = info.borrow_mut().new_task_id();
let handle = JoinHandle::new(info.clone(), id);
let task = Box::new(TaskImpl {
id,
stat: TaskStat::Init,
future,
handle: handle.weak(),
});
(task, handle)
}
pub(crate) fn with_join(
future: T,
info: &Rc<RefCell<SchedInfo>>,
handle: Weak<dyn Join<T::Output>>,
) -> Box<dyn Task> {
let id = info.borrow_mut().new_task_id();
Box::new(TaskImpl {
stat: TaskStat::Init,
id,
future,
handle,
})
}
fn set_running(&mut self) {
if let TaskStat::Init = self.stat {
self.stat = TaskStat::Running;
if let Some(handle) = self.handle.upgrade() {
handle.set_running(self.id);
}
}
}
}
impl<T: Future + 'static> Task for TaskImpl<T> {
fn get_id(&self) -> u64 {
self.id
}
fn abort(&mut self) {
self.stat = TaskStat::Aborted;
if let Some(handle) = self.handle.upgrade() {
handle.set_aborted(self.id);
}
}
fn run(&mut self, ctx: &mut Context<'_>) -> Poll<()> {
self.set_running();
let pin = unsafe { Pin::new_unchecked(&mut self.future) };
match pin.poll(ctx) {
Poll::Pending => Poll::Pending,
Poll::Ready(output) => {
if let Some(handle) = self.handle.upgrade() {
handle.set_finished(output, self.id);
}
self.stat = TaskStat::End;
Poll::Ready(())
}
}
}
}