Cache TaskLayout in Header.

This also allows debuggers to reliably decode tasks from raw memory blobs.

Author: michaelwoerister

src/header.rs

@@ -3,7 +3,7 @@ use core::fmt;
 use core::sync::atomic::{AtomicUsize, Ordering};
 use core::task::Waker;
 
-use crate::raw::TaskVTable;
+use crate::raw::{TaskLayout, TaskVTable};
 use crate::state::*;
 use crate::utils::abort_on_panic;
 
@@ -26,6 +26,10 @@ pub(crate) struct Header {
     /// In addition to the actual waker virtual table, it also contains pointers to several other
     /// methods necessary for bookkeeping the heap-allocated task.
     pub(crate) vtable: &'static TaskVTable,
+
+    /// The memory layout of the task. Caching this in the header speeds things up slightly,
+    /// and it also enables debuggers to decode raw task memory blobs.
+    pub(crate) layout: TaskLayout,
 }
 
 impl Header {

src/raw.rs

@@ -1,5 +1,6 @@
 use alloc::alloc::Layout;
 use core::cell::UnsafeCell;
+use core::convert::TryInto;
 use core::future::Future;
 use core::mem::{self, ManuallyDrop};
 use core::pin::Pin;
@@ -45,16 +46,22 @@ pub(crate) struct TaskVTable {
 #[derive(Clone, Copy)]
 pub(crate) struct TaskLayout {
     /// Memory layout of the whole task.
-    pub(crate) layout: Layout,
+    pub(crate) size: u32,
+    pub(crate) align: u32,
 
     /// Offset into the task at which the schedule function is stored.
-    pub(crate) offset_s: usize,
+    pub(crate) offset_schedule: u32,
 
-    /// Offset into the task at which the future is stored.
-    pub(crate) offset_f: usize,
+    /// Offset into the task at which the future or the output is stored.
+    pub(crate) offset_future_or_output: u32,
+}
 
-    /// Offset into the task at which the output is stored.
-    pub(crate) offset_r: usize,
+impl TaskLayout {
+    #[inline]
+    fn layout(self) -> Layout {
+        // This is safe because `size` and `align` always come for a valid `Layout` value.
+        unsafe { Layout::from_size_align_unchecked(self.size as usize, self.align as usize) }
+    }
 }
 
 /// Raw pointers to the fields inside a task.
@@ -101,15 +108,14 @@ where
 
         unsafe {
             // Allocate enough space for the entire task.
-            let ptr = match NonNull::new(alloc::alloc::alloc(task_layout.layout) as *mut ()) {
+            let ptr = match NonNull::new(alloc::alloc::alloc(task_layout.layout()) as *mut ()) {
                 None => abort(),
                 Some(p) => p,
             };
 
-            let raw = Self::from_ptr(ptr.as_ptr());
-
-            // Write the header as the first field of the task.
-            (raw.header as *mut Header).write(Header {
+            // Write the header as the first field of the task. The header must be initialized
+            // before `Self::from_ptr()` can be called.
+            (ptr.as_ptr() as *mut Header).write(Header {
                 state: AtomicUsize::new(SCHEDULED | TASK | REFERENCE),
                 awaiter: UnsafeCell::new(None),
                 vtable: &TaskVTable {
@@ -121,8 +127,11 @@ where
                     run: Self::run,
                     clone_waker: Self::clone_waker,
                 },
+                layout: task_layout,
             });
 
+            let raw = Self::from_ptr(ptr.as_ptr());
+
             // Write the schedule function as the third field of the task.
             (raw.schedule as *mut S).write(schedule);
 
@@ -136,15 +145,18 @@ where
     /// Creates a `RawTask` from a raw task pointer.
     #[inline]
     pub(crate) fn from_ptr(ptr: *const ()) -> Self {
-        let task_layout = Self::task_layout();
         let p = ptr as *const u8;
+        let p_header = p as *const Header;
 
         unsafe {
+            let offset_schedule = (*p_header).layout.offset_schedule as usize;
+            let offset_future_or_output = (*p_header).layout.offset_future_or_output as usize;
+
             Self {
-                header: p as *const Header,
-                schedule: p.add(task_layout.offset_s) as *const S,
-                future: p.add(task_layout.offset_f) as *mut F,
-                output: p.add(task_layout.offset_r) as *mut T,
+                header: p_header,
+                schedule: p.add(offset_schedule) as *const S,
+                future: p.add(offset_future_or_output) as *mut F,
+                output: p.add(offset_future_or_output) as *mut T,
             }
         }
     }
@@ -165,16 +177,20 @@ where
 
         // Compute the layout for `Header` followed `S` and `union { F, T }`.
         let layout = layout_header;
-        let (layout, offset_s) = extend(layout, layout_s);
+        let (layout, offset_schedule) = extend(layout, layout_s);
         let (layout, offset_union) = extend(layout, layout_union);
-        let offset_f = offset_union;
-        let offset_r = offset_union;
+
+        // Converting to `u32` is unlikely to fail.
+        let size: u32 = layout.size().try_into().unwrap();
+        let align: u32 = layout.align().try_into().unwrap();
+        let offset_schedule: u32 = offset_schedule.try_into().unwrap();
+        let offset_future_or_output: u32 = offset_union.try_into().unwrap();
 
         TaskLayout {
-            layout,
-            offset_s,
-            offset_f,
-            offset_r,
+            size,
+            align,
+            offset_schedule,
+            offset_future_or_output,
         }
     }
 
@@ -416,7 +432,6 @@ where
     #[inline]
     unsafe fn destroy(ptr: *const ()) {
         let raw = Self::from_ptr(ptr);
-        let task_layout = Self::task_layout();
 
         // We need a safeguard against panics because destructors can panic.
         abort_on_panic(|| {
@@ -425,7 +440,7 @@ where
         });
 
         // Finally, deallocate the memory reserved by the task.
-        alloc::alloc::dealloc(ptr as *mut u8, task_layout.layout);
+        alloc::alloc::dealloc(ptr as *mut u8, (*raw.header).layout.layout());
     }
 
     /// Runs a task.