monoio/task/
raw.rs

1use std::{
2    future::Future,
3    ptr::NonNull,
4    task::{Poll, Waker},
5};
6
7use crate::task::{Cell, Harness, Header, Schedule};
8
9pub(crate) struct RawTask {
10    ptr: NonNull<Header>,
11}
12
13impl Clone for RawTask {
14    fn clone(&self) -> Self {
15        *self
16    }
17}
18
19impl Copy for RawTask {}
20
21pub(crate) struct Vtable {
22    /// Poll the future
23    pub(crate) poll: unsafe fn(NonNull<Header>),
24    /// Deallocate the memory
25    pub(crate) dealloc: unsafe fn(NonNull<Header>),
26
27    /// Read the task output, if complete
28    pub(crate) try_read_output: unsafe fn(NonNull<Header>, *mut (), &Waker),
29
30    /// The join handle has been dropped
31    pub(crate) drop_join_handle_slow: unsafe fn(NonNull<Header>),
32
33    /// Set future output
34    #[cfg(feature = "sync")]
35    pub(crate) finish: unsafe fn(NonNull<Header>, *mut ()),
36}
37
38/// Get the vtable for the requested `T` and `S` generics.
39pub(super) fn vtable<T: Future, S: Schedule>() -> &'static Vtable {
40    &Vtable {
41        poll: poll::<T, S>,
42        dealloc: dealloc::<T, S>,
43        try_read_output: try_read_output::<T, S>,
44        drop_join_handle_slow: drop_join_handle_slow::<T, S>,
45        #[cfg(feature = "sync")]
46        finish: finish::<T, S>,
47    }
48}
49
50impl RawTask {
51    pub(crate) fn new<T, S>(owner_id: usize, task: T, scheduler: S) -> RawTask
52    where
53        T: Future,
54        S: Schedule,
55    {
56        let ptr = Box::into_raw(Cell::new(owner_id, task, scheduler));
57        let ptr = unsafe { NonNull::new_unchecked(ptr as *mut Header) };
58
59        RawTask { ptr }
60    }
61
62    pub(crate) unsafe fn from_raw(ptr: NonNull<Header>) -> RawTask {
63        RawTask { ptr }
64    }
65
66    pub(crate) fn header(&self) -> &Header {
67        unsafe { self.ptr.as_ref() }
68    }
69
70    /// Safety: mutual exclusion is required to call this function.
71    pub(crate) fn poll(self) {
72        let vtable = self.header().vtable;
73        unsafe { (vtable.poll)(self.ptr) }
74    }
75
76    pub(crate) fn dealloc(self) {
77        let vtable = self.header().vtable;
78        unsafe {
79            (vtable.dealloc)(self.ptr);
80        }
81    }
82
83    /// Safety: `dst` must be a `*mut Poll<super::Result<T::Output>>` where `T`
84    /// is the future stored by the task.
85    pub(crate) unsafe fn try_read_output(self, dst: *mut (), waker: &Waker) {
86        let vtable = self.header().vtable;
87        (vtable.try_read_output)(self.ptr, dst, waker);
88    }
89
90    pub(crate) fn drop_join_handle_slow(self) {
91        let vtable = self.header().vtable;
92        unsafe { (vtable.drop_join_handle_slow)(self.ptr) }
93    }
94
95    #[cfg(feature = "sync")]
96    pub(crate) unsafe fn finish(self, val_slot: *mut ()) {
97        let vtable = self.header().vtable;
98        unsafe { (vtable.finish)(self.ptr, val_slot) }
99    }
100}
101
102unsafe fn poll<T: Future, S: Schedule>(ptr: NonNull<Header>) {
103    let harness = Harness::<T, S>::from_raw(ptr);
104    harness.poll();
105}
106
107unsafe fn dealloc<T: Future, S: Schedule>(ptr: NonNull<Header>) {
108    let harness = Harness::<T, S>::from_raw(ptr);
109    harness.dealloc();
110}
111
112#[cfg(feature = "sync")]
113unsafe fn finish<T: Future, S: Schedule>(ptr: NonNull<Header>, val: *mut ()) {
114    let harness = Harness::<T, S>::from_raw(ptr);
115    let val = &mut *(val as *mut Option<<T as Future>::Output>);
116    harness.finish(val.take().unwrap());
117}
118
119unsafe fn try_read_output<T: Future, S: Schedule>(
120    ptr: NonNull<Header>,
121    dst: *mut (),
122    waker: &Waker,
123) {
124    let out = &mut *(dst as *mut Poll<T::Output>);
125
126    let harness = Harness::<T, S>::from_raw(ptr);
127    harness.try_read_output(out, waker);
128}
129
130unsafe fn drop_join_handle_slow<T: Future, S: Schedule>(ptr: NonNull<Header>) {
131    let harness = Harness::<T, S>::from_raw(ptr);
132    harness.drop_join_handle_slow()
133}
monoio/task/raw.rs

monoio/task/
raw.rs
