rquickjs_core/runtime/schedular/
task.rs

1use std::{
2    cell::{Cell, UnsafeCell},
3    future::Future,
4    mem::ManuallyDrop,
5    ptr::{addr_of_mut, NonNull},
6    sync::{atomic::AtomicBool, Arc, Weak},
7    task::{Context, Poll},
8};
9
10use super::{
11    queue::{NodeHeader, Queue},
12    vtable::VTable,
13};
14
15#[repr(C)]
16pub struct Task<F> {
17    /// Header for the intrusive list,
18    /// Must be first.
19    pub(crate) head: NodeHeader,
20    /// Data the schedular uses to run the future.
21    pub(crate) body: TaskBody,
22    /// The future itself.
23    pub(crate) future: UnsafeCell<ManuallyDrop<F>>,
24}
25
26impl<F: Future<Output = ()>> Task<F> {
27    pub fn new(queue: Weak<Queue>, f: F) -> Self {
28        Self {
29            head: NodeHeader::new(),
30            body: TaskBody {
31                queue,
32                vtable: VTable::get::<F>(),
33                next: Cell::new(None),
34                prev: Cell::new(None),
35                queued: AtomicBool::new(true),
36                done: Cell::new(false),
37            },
38            future: UnsafeCell::new(ManuallyDrop::new(f)),
39        }
40    }
41}
42
43// Seperate struct to not have everything in task be repr(C)
44pub struct TaskBody {
45    pub(crate) queue: Weak<Queue>,
46    pub(crate) vtable: &'static VTable,
47    // The double linked list of tasks.
48    pub(crate) next: Cell<Option<ErasedTaskPtr>>,
49    pub(crate) prev: Cell<Option<ErasedTaskPtr>>,
50    // wether the task is currently in the queue to be re-polled.
51    pub(crate) queued: AtomicBool,
52    pub(crate) done: Cell<bool>,
53}
54
55/// A raw pointer to a task with it's type erased.
56#[derive(Clone, Copy, Debug)]
57#[repr(transparent)]
58pub struct ErasedTaskPtr(NonNull<Task<()>>);
59
60impl ErasedTaskPtr {
61    pub fn from_nonnull(ptr: NonNull<Task<()>>) -> Self {
62        Self(ptr)
63    }
64
65    pub unsafe fn body<'a>(self) -> &'a TaskBody {
66        let ptr = self.0.as_ptr();
67        unsafe { &*addr_of_mut!((*ptr).body) }
68    }
69
70    pub fn as_node_ptr(self) -> NonNull<NodeHeader> {
71        self.0.cast()
72    }
73
74    pub fn as_nonnull(self) -> NonNull<Task<()>> {
75        self.0
76    }
77
78    pub unsafe fn task_drive(self, cx: &mut Context) -> Poll<()> {
79        unsafe { (self.body().vtable.task_drive)(self.0, cx) }
80    }
81
82    pub unsafe fn task_incr(self) {
83        unsafe { (self.body().vtable.task_incr)(self.0) }
84    }
85
86    pub unsafe fn task_decr(self) {
87        unsafe { (self.body().vtable.task_decr)(self.0) }
88    }
89
90    pub unsafe fn task_drop(self) {
91        unsafe { (self.body().vtable.task_drop)(self.0) }
92    }
93}
94
95/// An owning pointer to a task with it's type erased.
96pub struct ErasedTask(ErasedTaskPtr);
97
98impl ErasedTask {
99    pub unsafe fn from_ptr(ptr: ErasedTaskPtr) -> Self {
100        Self(ptr)
101    }
102
103    pub fn into_ptr(this: Self) -> ErasedTaskPtr {
104        let res = this.0;
105        std::mem::forget(this);
106        res
107    }
108
109    pub fn new<F>(task: Arc<Task<F>>) -> Self {
110        unsafe {
111            let ptr = NonNull::new_unchecked(Arc::into_raw(task) as *mut Task<F>).cast();
112            Self(ErasedTaskPtr(ptr))
113        }
114    }
115
116    pub fn body(&self) -> &TaskBody {
117        unsafe { self.0.body() }
118    }
119}
120
121impl Clone for ErasedTask {
122    fn clone(&self) -> Self {
123        unsafe {
124            self.0.task_incr();
125            Self(self.0)
126        }
127    }
128}
129
130impl Drop for ErasedTask {
131    fn drop(&mut self) {
132        unsafe { self.0.task_decr() }
133    }
134}