use super::{status::*, Attr, Extensions, RawTask, RawTaskVTable, TaskRef};
use crate::thread::Semaphore;
use crate::{err, Error, Result, CacheLineAligned};
use core::alloc::Layout;
use core::cell::Cell;
use core::future::Future;
use core::mem::ManuallyDrop;
use core::pin::Pin;
use core::ptr::{self, NonNull};
use core::result;
use core::sync::atomic::{
fence,
Ordering::{self, Acquire, Relaxed, Release},
};
use core::task::{Context, Poll, Waker};
use core::time::Duration;
use hipool::{Boxed, MemPool, NullAlloc};
enum Body<T: Future> {
Running(T),
Return(T::Output),
Clean,
}
enum Joinable {
Sync(NonNull<Semaphore>),
Async(Waker),
Null,
}
#[repr(C)]
pub(crate) struct Task<T: Future> {
base: RawTask,
body: Cell<Body<T>>,
joinable: Cell<Joinable>,
pool: ManuallyDrop<Boxed<'static, MemPool, NullAlloc>>,
}
impl<T: Future> Task<T> {
#[allow(clippy::new_ret_no_self)]
pub fn new(future: T, attr: &Attr) -> Result<TaskRef> {
let pool = MemPool::new_boxed(0)?;
let (pool, layout, alloc) = pool.leak();
let boxed = unsafe { Boxed::from_with(pool.into(), layout, alloc) };
Self::new_in(pool, boxed, future, attr)
}
}
impl<T: Future> Task<T> {
const TASK_VTABLE: RawTaskVTable = RawTaskVTable {
release: Self::release,
poll: Self::poll,
exit: Self::exit,
output: Self::output,
join: Self::join,
tail: Self::tail,
};
fn new_in(
pool: &'static MemPool,
boxed: Boxed<'static, MemPool, NullAlloc>,
future: T,
attr: &Attr,
) -> Result<TaskRef> {
if attr.timeout == Duration::MAX {
Task::new_with(pool, boxed, future, attr)
} else if attr.delay {
Task::new_with(pool, boxed, future.delay(attr.timeout), attr)
} else {
Task::new_with(pool, boxed, future.deadline(attr.timeout), attr)
}
}
fn new_with(
pool: &'static MemPool,
boxed: Boxed<'static, MemPool, NullAlloc>,
future: T,
attr: &Attr,
) -> Result<TaskRef> {
let task = Boxed::new_in(
pool,
CacheLineAligned::new(Self {
base: RawTask::new(&Self::TASK_VTABLE, attr),
body: Cell::new(Body::Running(future)),
joinable: Cell::new(Joinable::Null),
pool: ManuallyDrop::new(boxed),
}),
)?
.leak()
.0;
Ok(TaskRef::new(NonNull::from(&mut task.base)))
}
unsafe fn from_raw(task: &RawTask) -> &Self {
hioff::container_of!(task, Self, base)
}
unsafe fn from_raw_mut(task: &mut RawTask) -> *mut Self {
hioff::container_of_mut!(task, Self, base)
}
fn exit_with(&mut self, task: &RawTask, ret: result::Result<T::Output, u32>) {
let exit_code = match ret {
Ok(ret) => {
self.body.set(Body::Return(ret));
STAT_FINISH
}
Err(ret) => ret,
};
let status = task.status.cmp_xchg_status(STAT_RUN, exit_code, Release);
match status {
Ok(_status) => {}
Err(STAT_RETURN) => {
task.status.set_status(exit_code, Release);
fence(Ordering::Acquire);
match unsafe { &*self.joinable.as_ptr() } {
Joinable::Async(waker) => waker.wake_by_ref(),
_ => unreachable!("expect Joinable::Async"),
}
}
Err(STAT_JOIN) => unsafe {
task.status.set_status(exit_code, Release);
fence(Ordering::Acquire);
match self.joinable.replace(Joinable::Null) {
Joinable::Sync(mut sem) => sem.as_mut().post(),
_ => unreachable!("expect Joinable::Sync"),
}
},
Err(_status) => {}
}
}
fn poll(task: &mut RawTask, ctx: &mut Context<'_>) -> Poll<()> {
let this = unsafe { &mut *Self::from_raw_mut(task) };
let future = match this.body.get_mut() {
Body::Running(future) => future,
_ => unreachable!("error body, should be Running"),
};
let pinned = unsafe { Pin::new_unchecked(future) };
match Future::poll(pinned, ctx) {
Poll::Pending => Poll::Pending,
Poll::Ready(ret) => {
this.exit_with(task, Ok(ret));
Poll::Ready(())
}
}
}
fn exit(task: &mut RawTask, stat: u32) {
let this = unsafe { &mut *Self::from_raw_mut(task) };
this.exit_with(task, Err(stat));
}
unsafe fn release(task: &mut RawTask) {
let this = unsafe { &mut *Self::from_raw_mut(task) };
let _pool = unsafe { ManuallyDrop::take(&mut this.pool) };
unsafe { ptr::drop_in_place(this) };
unsafe { ptr::drop_in_place(task) };
}
unsafe fn exit_code(&self, task: &RawTask, output: *mut ()) -> Result<()> {
match task.status.status(Relaxed) {
STAT_FINISH => {
let output = output.cast::<T::Output>();
match self.body.replace(Body::Clean) {
Body::Return(val) => output.write(val),
_ => unreachable!("error body, should be Return"),
}
Ok(())
}
STAT_EXIT => {
self.body.replace(Body::Clean);
Err(Error::new(err::ESHUTDOWN))
}
STAT_ABORT => {
self.body.replace(Body::Clean);
Err(Error::new(err::ECANCELED))
}
_ => {
self.body.replace(Body::Clean);
Err(Error::default())
}
}
}
unsafe fn join(task: &RawTask, output: *mut ()) -> Result<()> {
let this = unsafe { Self::from_raw(task) };
let sem = Semaphore::new()?;
this.joinable.set(Joinable::Sync(NonNull::from(&sem)));
let status = task.status.cmp_xchg_status(STAT_RUN, STAT_JOIN, Release);
if status.is_ok() {
sem.wait();
} else {
fence(Acquire);
}
this.joinable.set(Joinable::Null);
this.exit_code(task, output)
}
unsafe fn output(task: &RawTask, waker: &Waker, output: *mut ()) -> Poll<Result<()>> {
let this = unsafe { Self::from_raw(task) };
let joined = match unsafe { &*this.joinable.as_ptr() } {
Joinable::Async(_) => true,
Joinable::Null => false,
_ => unreachable!("Expect Joinable::Async or Joinable::Null"),
};
if !joined {
this.joinable.set(Joinable::Async(waker.clone()));
let _ = task.status.cmp_xchg_status(STAT_RUN, STAT_RETURN, Release);
}
if task.status.status(Acquire) == STAT_RETURN {
Poll::Pending
} else {
this.joinable.set(Joinable::Null);
Poll::Ready(this.exit_code(task, output))
}
}
unsafe fn tail(task: &RawTask, layout: Layout) -> *const () {
let this = unsafe { Self::from_raw(task) } as *const Self;
let offset = Layout::new::<Self>().extend(layout).unwrap().1;
this.cast::<u8>().add(offset).cast::<()>()
}
}