use super::{status::*, Attr, TaskContext, TaskRef, TaskWaker, Worker};
use crate::event::{Event, Scheduler};
use crate::Result;
use core::alloc::Layout;
use core::mem::ManuallyDrop;
use core::ptr::{self, NonNull};
use core::sync::atomic::Ordering::Relaxed;
use core::task::{Context, Poll, Waker};
use hicollections::List;
use hioff::container_of_mut;
pub(crate) union Node {
next: *mut RawTask,
event: ManuallyDrop<Event>,
}
#[cfg(feature = "task_mem_split")]
#[repr(C)]
pub struct RawTask {
pub(crate) node: Node,
pub(crate) vtable: &'static RawTaskVTable,
pub(crate) status: AtomicStatus,
pub(crate) private: *const (),
pub(crate) future: *const (),
}
#[cfg(not(feature = "task_mem_split"))]
#[repr(C)]
pub struct RawTask {
pub(crate) node: Node,
pub(crate) vtable: &'static RawTaskVTable,
pub(crate) status: AtomicStatus,
pub(crate) private: *const (),
}
impl Node {
fn new() -> Self {
Self {
next: ptr::null_mut::<RawTask>(),
}
}
pub(crate) unsafe fn set_next(&mut self, task: *mut RawTask) {
self.next = task;
}
pub(crate) unsafe fn init_next(&mut self) {
self.next = ptr::null_mut::<RawTask>();
}
pub(crate) unsafe fn next(&self) -> *mut RawTask {
self.next
}
pub(crate) unsafe fn init_event(
&mut self,
handle: fn(&Event, u32, &mut Scheduler),
) -> &mut Event {
self.event = ManuallyDrop::new(Event::new(handle));
&mut self.event
}
}
impl RawTask {
#[cfg(feature = "task_mem_split")]
pub(crate) fn new(vtable: &'static RawTaskVTable, attr: &Attr, future: *const ()) -> Self {
Self {
node: Node::new(),
vtable,
status: AtomicStatus::new(attr),
private: ptr::null(),
future,
}
}
#[cfg(not(feature = "task_mem_split"))]
pub(crate) fn new(vtable: &'static RawTaskVTable, attr: &Attr) -> Self {
Self {
node: Node::new(),
vtable,
status: AtomicStatus::new(attr),
private: ptr::null(),
}
}
pub(crate) unsafe fn from_event(e: &Event) -> *mut Self {
container_of_mut!(e, Self, node.event)
}
pub(crate) fn inc_ref(&self) {
self.status.inc_ref();
}
pub(crate) fn dec_ref(&self) {
if self.status.test_dec_ref() {
let this = unsafe { core::ptr::NonNull::from(self).as_mut() };
unsafe { this.release() };
}
}
pub(crate) fn task_ref(&self) -> TaskRef {
self.inc_ref();
TaskRef::new(NonNull::from(self))
}
}
impl TaskRef {
pub(crate) fn fast_wake(self, sched: &mut Scheduler) {
if self.status.test_inc_wake() {
self.sched_waked(sched);
}
}
pub(crate) fn sched_waked(self, sched: &mut Scheduler) {
let mut task = ManuallyDrop::new(self);
let pri = task.status.priority() as i32;
let e = unsafe { task.node.init_event(RawTask::event_handle) };
unsafe { sched.add_event(e, 0, pri) };
}
pub(crate) fn add_waked(self, events: &mut List<Event>) {
let mut task = ManuallyDrop::new(self);
let e = unsafe { task.node.init_event(RawTask::event_handle) };
unsafe { events.add_tail(e) };
}
}
impl RawTask {
pub(crate) fn is_finished(&self) -> bool {
let status = self.status.status(Relaxed);
status >= STAT_FINISH
}
pub(crate) fn abort(&self) -> bool {
if self
.status
.cmp_xchg_status(STAT_RUN, STAT_ABORT, Relaxed)
.is_ok()
{
self.wake();
return true;
}
false
}
pub(crate) fn wake(&self) {
if self.status.test_inc_wake() {
let group_id = self.status.group();
let local = self.status.get_local();
if let Some(current) = Worker::current() {
if current.group_id() == group_id {
return current.wake(self.task_ref(), local);
}
}
super::Runtime::get(group_id).sched(self.task_ref(), local);
}
}
pub(crate) fn run(&mut self, sched: &mut Scheduler) {
let addr = sched as *const _ as *const u8 as usize;
let task_waker = TaskWaker::new(self, sched);
let mut ctx = Context::from_waker(task_waker.waker());
while self.run_continue(addr) {
let cnt = self.status.wake_count();
self.status.poll_inc();
if self.poll(&mut ctx) == Poll::Pending && self.status.test_dec_wake(cnt) {
continue;
}
if self.status.exited() {
break;
}
return;
}
if self.status.poll_cnt() > 0 && !self.status.aborted() {
self.status.set_aborted();
ctx.set_aborted(true);
let _ = self.poll(&mut ctx);
}
}
#[inline(always)]
fn run_continue(&self, _sched: usize) -> bool {
let status = self.status.status(Relaxed);
if status < STAT_FINISH {
true
} else {
assert!(status == STAT_ABORT || status == STAT_EXIT);
false
}
}
fn event_handle(event: &Event, _events: u32, sched: &mut Scheduler) {
let this = unsafe { &mut *Self::from_event(event) };
this.run(sched);
this.dec_ref();
}
}
#[repr(C)]
pub struct RawTaskVTable {
pub poll: fn(&mut RawTask, ctx: &mut Context<'_>) -> Poll<()>,
pub release: unsafe fn(&mut RawTask),
pub output: unsafe fn(&RawTask, waker: &Waker, output: *mut ()) -> Poll<Result<()>>,
pub join: unsafe fn(&RawTask, output: *mut ()) -> Result<()>,
pub exit: fn(&mut RawTask),
pub tail: unsafe fn(&RawTask, layout: Layout) -> *const (),
}
impl RawTask {
unsafe fn release(&mut self) {
(self.vtable.release)(self)
}
pub fn poll(&mut self, ctx: &mut Context<'_>) -> Poll<()> {
(self.vtable.poll)(self, ctx)
}
pub fn exit(&mut self) {
(self.vtable.exit)(self)
}
pub unsafe fn output(&self, waker: &Waker, output: *mut ()) -> Poll<Result<()>> {
(self.vtable.output)(self, waker, output)
}
pub unsafe fn join(&self, output: *mut ()) -> Result<()> {
(self.vtable.join)(self, output)
}
pub unsafe fn tail(&self, layout: Layout) -> *const () {
(self.vtable.tail)(self, layout)
}
}