use crate::{Join, JoinHandle, SchedInfo};
use std::cell::RefCell;
use std::future::Future;
use std::pin::Pin;
use std::rc::{Rc, Weak};
use std::task::{Context, Poll};
#[derive(Copy, Clone)]
pub(crate) enum TaskState {
Pending,
Running,
}
pub(crate) trait Task {
fn get_id(&self) -> u64;
fn get_status(&self) -> TaskState;
fn set_running(&mut self);
fn run(&mut self, ctx: &mut Context<'_>) -> Poll<()>;
}
pub(crate) struct TaskImpl<T: Future> {
id: u64,
stat: TaskState,
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>) {
let id = info.borrow_mut().new_task_id();
let handle = JoinHandle::new(info.clone(), id);
let task = Box::new(TaskImpl {
id,
stat: TaskState::Pending,
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: TaskState::Pending,
id,
future,
handle,
})
}
}
impl<T: Future> Task for TaskImpl<T> {
fn get_id(&self) -> u64 {
self.id
}
fn get_status(&self) -> TaskState {
self.stat
}
fn set_running(&mut self) {
self.stat = TaskState::Running;
if let Some(handle) = self.handle.upgrade() {
handle.set_running(self.id);
}
}
fn run(&mut self, ctx: &mut Context<'_>) -> Poll<()> {
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);
}
Poll::Ready(())
}
}
}
}