tokio/runtime/task/
raw.rs

1use crate::future::Future;
2use crate::runtime::task::core::{Core, Trailer};
3use crate::runtime::task::{Cell, Harness, Header, Id, Schedule, State};
4#[cfg(tokio_unstable)]
5use std::panic::Location;
6use std::ptr::NonNull;
7use std::task::{Poll, Waker};
8
9/// Raw task handle
10#[derive(Clone)]
11pub(crate) struct RawTask {
12    ptr: NonNull<Header>,
13}
14
15pub(super) struct Vtable {
16    /// Polls the future.
17    pub(super) poll: unsafe fn(NonNull<Header>),
18
19    /// Schedules the task for execution on the runtime.
20    pub(super) schedule: unsafe fn(NonNull<Header>),
21
22    /// Deallocates the memory.
23    pub(super) dealloc: unsafe fn(NonNull<Header>),
24
25    /// Reads the task output, if complete.
26    pub(super) try_read_output: unsafe fn(NonNull<Header>, *mut (), &Waker),
27
28    /// The join handle has been dropped.
29    pub(super) drop_join_handle_slow: unsafe fn(NonNull<Header>),
30
31    /// An abort handle has been dropped.
32    pub(super) drop_abort_handle: unsafe fn(NonNull<Header>),
33
34    /// Scheduler is being shutdown.
35    pub(super) shutdown: unsafe fn(NonNull<Header>),
36
37    /// The number of bytes that the `trailer` field is offset from the header.
38    pub(super) trailer_offset: usize,
39
40    /// The number of bytes that the `scheduler` field is offset from the header.
41    pub(super) scheduler_offset: usize,
42
43    /// The number of bytes that the `id` field is offset from the header.
44    pub(super) id_offset: usize,
45
46    /// The number of bytes that the `spawned_at` field is offset from the header.
47    #[cfg(tokio_unstable)]
48    pub(super) spawn_location_offset: usize,
49}
50
51/// Get the vtable for the requested `T` and `S` generics.
52pub(super) fn vtable<T: Future, S: Schedule>() -> &'static Vtable {
53    &Vtable {
54        poll: poll::<T, S>,
55        schedule: schedule::<S>,
56        dealloc: dealloc::<T, S>,
57        try_read_output: try_read_output::<T, S>,
58        drop_join_handle_slow: drop_join_handle_slow::<T, S>,
59        drop_abort_handle: drop_abort_handle::<T, S>,
60        shutdown: shutdown::<T, S>,
61        trailer_offset: OffsetHelper::<T, S>::TRAILER_OFFSET,
62        scheduler_offset: OffsetHelper::<T, S>::SCHEDULER_OFFSET,
63        id_offset: OffsetHelper::<T, S>::ID_OFFSET,
64        #[cfg(tokio_unstable)]
65        spawn_location_offset: OffsetHelper::<T, S>::SPAWN_LOCATION_OFFSET,
66    }
67}
68
69/// Calling `get_trailer_offset` directly in vtable doesn't work because it
70/// prevents the vtable from being promoted to a static reference.
71///
72/// See this thread for more info:
73/// <https://users.rust-lang.org/t/custom-vtables-with-integers/78508>
74struct OffsetHelper<T, S>(T, S);
75impl<T: Future, S: Schedule> OffsetHelper<T, S> {
76    // Pass `size_of`/`align_of` as arguments rather than calling them directly
77    // inside `get_trailer_offset` because trait bounds on generic parameters
78    // of const fn are unstable on our MSRV.
79    const TRAILER_OFFSET: usize = get_trailer_offset(
80        std::mem::size_of::<Header>(),
81        std::mem::size_of::<Core<T, S>>(),
82        std::mem::align_of::<Core<T, S>>(),
83        std::mem::align_of::<Trailer>(),
84    );
85
86    // The `scheduler` is the first field of `Core`, so it has the same
87    // offset as `Core`.
88    const SCHEDULER_OFFSET: usize = get_core_offset(
89        std::mem::size_of::<Header>(),
90        std::mem::align_of::<Core<T, S>>(),
91    );
92
93    const ID_OFFSET: usize = get_id_offset(
94        std::mem::size_of::<Header>(),
95        std::mem::align_of::<Core<T, S>>(),
96        std::mem::size_of::<S>(),
97        std::mem::align_of::<Id>(),
98    );
99
100    #[cfg(tokio_unstable)]
101    const SPAWN_LOCATION_OFFSET: usize = get_spawn_location_offset(
102        std::mem::size_of::<Header>(),
103        std::mem::align_of::<Core<T, S>>(),
104        std::mem::size_of::<S>(),
105        std::mem::align_of::<Id>(),
106        std::mem::size_of::<Id>(),
107        std::mem::align_of::<&'static Location<'static>>(),
108    );
109}
110
111/// Compute the offset of the `Trailer` field in `Cell<T, S>` using the
112/// `#[repr(C)]` algorithm.
113///
114/// Pseudo-code for the `#[repr(C)]` algorithm can be found here:
115/// <https://doc.rust-lang.org/reference/type-layout.html#reprc-structs>
116const fn get_trailer_offset(
117    header_size: usize,
118    core_size: usize,
119    core_align: usize,
120    trailer_align: usize,
121) -> usize {
122    let mut offset = header_size;
123
124    let core_misalign = offset % core_align;
125    if core_misalign > 0 {
126        offset += core_align - core_misalign;
127    }
128    offset += core_size;
129
130    let trailer_misalign = offset % trailer_align;
131    if trailer_misalign > 0 {
132        offset += trailer_align - trailer_misalign;
133    }
134
135    offset
136}
137
138/// Compute the offset of the `Core<T, S>` field in `Cell<T, S>` using the
139/// `#[repr(C)]` algorithm.
140///
141/// Pseudo-code for the `#[repr(C)]` algorithm can be found here:
142/// <https://doc.rust-lang.org/reference/type-layout.html#reprc-structs>
143const fn get_core_offset(header_size: usize, core_align: usize) -> usize {
144    let mut offset = header_size;
145
146    let core_misalign = offset % core_align;
147    if core_misalign > 0 {
148        offset += core_align - core_misalign;
149    }
150
151    offset
152}
153
154/// Compute the offset of the `Id` field in `Cell<T, S>` using the
155/// `#[repr(C)]` algorithm.
156///
157/// Pseudo-code for the `#[repr(C)]` algorithm can be found here:
158/// <https://doc.rust-lang.org/reference/type-layout.html#reprc-structs>
159const fn get_id_offset(
160    header_size: usize,
161    core_align: usize,
162    scheduler_size: usize,
163    id_align: usize,
164) -> usize {
165    let mut offset = get_core_offset(header_size, core_align);
166    offset += scheduler_size;
167
168    let id_misalign = offset % id_align;
169    if id_misalign > 0 {
170        offset += id_align - id_misalign;
171    }
172
173    offset
174}
175
176/// Compute the offset of the `&'static Location<'static>` field in `Cell<T, S>`
177/// using the `#[repr(C)]` algorithm.
178///
179/// Pseudo-code for the `#[repr(C)]` algorithm can be found here:
180/// <https://doc.rust-lang.org/reference/type-layout.html#reprc-structs>
181#[cfg(tokio_unstable)]
182const fn get_spawn_location_offset(
183    header_size: usize,
184    core_align: usize,
185    scheduler_size: usize,
186    id_align: usize,
187    id_size: usize,
188    spawn_location_align: usize,
189) -> usize {
190    let mut offset = get_id_offset(header_size, core_align, scheduler_size, id_align);
191    offset += id_size;
192
193    let spawn_location_misalign = offset % spawn_location_align;
194    if spawn_location_misalign > 0 {
195        offset += spawn_location_align - spawn_location_misalign;
196    }
197
198    offset
199}
200
201impl RawTask {
202    pub(super) fn new<T, S>(
203        task: T,
204        scheduler: S,
205        id: Id,
206        _spawned_at: super::SpawnLocation,
207    ) -> RawTask
208    where
209        T: Future,
210        S: Schedule,
211    {
212        let ptr = Box::into_raw(Cell::<_, S>::new(
213            task,
214            scheduler,
215            State::new(),
216            id,
217            #[cfg(tokio_unstable)]
218            _spawned_at.0,
219        ));
220        let ptr = unsafe { NonNull::new_unchecked(ptr.cast()) };
221
222        RawTask { ptr }
223    }
224
225    pub(super) unsafe fn from_raw(ptr: NonNull<Header>) -> RawTask {
226        RawTask { ptr }
227    }
228
229    pub(super) fn header_ptr(&self) -> NonNull<Header> {
230        self.ptr
231    }
232
233    pub(super) fn trailer_ptr(&self) -> NonNull<Trailer> {
234        unsafe { Header::get_trailer(self.ptr) }
235    }
236
237    /// Returns a reference to the task's header.
238    pub(super) fn header(&self) -> &Header {
239        unsafe { self.ptr.as_ref() }
240    }
241
242    /// Returns a reference to the task's trailer.
243    pub(super) fn trailer(&self) -> &Trailer {
244        unsafe { &*self.trailer_ptr().as_ptr() }
245    }
246
247    /// Returns a reference to the task's state.
248    pub(super) fn state(&self) -> &State {
249        &self.header().state
250    }
251
252    /// Safety: mutual exclusion is required to call this function.
253    pub(crate) fn poll(self) {
254        let vtable = self.header().vtable;
255        unsafe { (vtable.poll)(self.ptr) }
256    }
257
258    pub(super) fn schedule(self) {
259        let vtable = self.header().vtable;
260        unsafe { (vtable.schedule)(self.ptr) }
261    }
262
263    pub(super) fn dealloc(self) {
264        let vtable = self.header().vtable;
265        unsafe {
266            (vtable.dealloc)(self.ptr);
267        }
268    }
269
270    /// Safety: `dst` must be a `*mut Poll<super::Result<T::Output>>` where `T`
271    /// is the future stored by the task.
272    pub(super) unsafe fn try_read_output(self, dst: *mut (), waker: &Waker) {
273        let vtable = self.header().vtable;
274        (vtable.try_read_output)(self.ptr, dst, waker);
275    }
276
277    pub(super) fn drop_join_handle_slow(self) {
278        let vtable = self.header().vtable;
279        unsafe { (vtable.drop_join_handle_slow)(self.ptr) }
280    }
281
282    pub(super) fn drop_abort_handle(self) {
283        let vtable = self.header().vtable;
284        unsafe { (vtable.drop_abort_handle)(self.ptr) }
285    }
286
287    pub(super) fn shutdown(self) {
288        let vtable = self.header().vtable;
289        unsafe { (vtable.shutdown)(self.ptr) }
290    }
291
292    /// Increment the task's reference count.
293    ///
294    /// Currently, this is used only when creating an `AbortHandle`.
295    pub(super) fn ref_inc(self) {
296        self.header().state.ref_inc();
297    }
298
299    /// Get the queue-next pointer
300    ///
301    /// This is for usage by the injection queue
302    ///
303    /// Safety: make sure only one queue uses this and access is synchronized.
304    pub(crate) unsafe fn get_queue_next(self) -> Option<RawTask> {
305        self.header()
306            .queue_next
307            .with(|ptr| *ptr)
308            .map(|p| RawTask::from_raw(p))
309    }
310
311    /// Sets the queue-next pointer
312    ///
313    /// This is for usage by the injection queue
314    ///
315    /// Safety: make sure only one queue uses this and access is synchronized.
316    pub(crate) unsafe fn set_queue_next(self, val: Option<RawTask>) {
317        self.header().set_next(val.map(|task| task.ptr));
318    }
319}
320
321impl Copy for RawTask {}
322
323unsafe fn poll<T: Future, S: Schedule>(ptr: NonNull<Header>) {
324    let harness = Harness::<T, S>::from_raw(ptr);
325    harness.poll();
326}
327
328unsafe fn schedule<S: Schedule>(ptr: NonNull<Header>) {
329    use crate::runtime::task::{Notified, Task};
330
331    let scheduler = Header::get_scheduler::<S>(ptr);
332    scheduler
333        .as_ref()
334        .schedule(Notified(Task::from_raw(ptr.cast())));
335}
336
337unsafe fn dealloc<T: Future, S: Schedule>(ptr: NonNull<Header>) {
338    let harness = Harness::<T, S>::from_raw(ptr);
339    harness.dealloc();
340}
341
342unsafe fn try_read_output<T: Future, S: Schedule>(
343    ptr: NonNull<Header>,
344    dst: *mut (),
345    waker: &Waker,
346) {
347    let out = &mut *(dst as *mut Poll<super::Result<T::Output>>);
348
349    let harness = Harness::<T, S>::from_raw(ptr);
350    harness.try_read_output(out, waker);
351}
352
353unsafe fn drop_join_handle_slow<T: Future, S: Schedule>(ptr: NonNull<Header>) {
354    let harness = Harness::<T, S>::from_raw(ptr);
355    harness.drop_join_handle_slow();
356}
357
358unsafe fn drop_abort_handle<T: Future, S: Schedule>(ptr: NonNull<Header>) {
359    let harness = Harness::<T, S>::from_raw(ptr);
360    harness.drop_reference();
361}
362
363unsafe fn shutdown<T: Future, S: Schedule>(ptr: NonNull<Header>) {
364    let harness = Harness::<T, S>::from_raw(ptr);
365    harness.shutdown();
366}
tokio/runtime/task/raw.rs

tokio/runtime/task/
raw.rs
