Initial commit
diff --git a/src/header.rs b/src/header.rs
new file mode 100644
index 0000000..0ce5164
--- /dev/null
+++ b/src/header.rs
@@ -0,0 +1,158 @@
+use std::alloc::Layout;
+use std::cell::Cell;
+use std::fmt;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::task::Waker;
+
+use crossbeam_utils::Backoff;
+
+use crate::raw::TaskVTable;
+use crate::state::*;
+use crate::utils::{abort_on_panic, extend};
+
+/// The header of a task.
+///
+/// This header is stored right at the beginning of every heap-allocated task.
+pub(crate) struct Header {
+ /// Current state of the task.
+ ///
+ /// Contains flags representing the current state and the reference count.
+ pub(crate) state: AtomicUsize,
+
+ /// The task that is blocked on the `JoinHandle`.
+ ///
+ /// This waker needs to be woken once the task completes or is closed.
+ pub(crate) awaiter: Cell<Option<Waker>>,
+
+ /// The virtual table.
+ ///
+ /// 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,
+}
+
+impl Header {
+ /// Cancels the task.
+ ///
+ /// This method will only mark the task as closed and will notify the awaiter, but it won't
+ /// reschedule the task if it's not completed.
+ pub(crate) fn cancel(&self) {
+ let mut state = self.state.load(Ordering::Acquire);
+
+ loop {
+ // If the task has been completed or closed, it can't be cancelled.
+ if state & (COMPLETED | CLOSED) != 0 {
+ break;
+ }
+
+ // Mark the task as closed.
+ match self.state.compare_exchange_weak(
+ state,
+ state | CLOSED,
+ Ordering::AcqRel,
+ Ordering::Acquire,
+ ) {
+ Ok(_) => {
+ // Notify the awaiter that the task has been closed.
+ if state & AWAITER != 0 {
+ self.notify();
+ }
+
+ break;
+ }
+ Err(s) => state = s,
+ }
+ }
+ }
+
+ /// Notifies the task blocked on the task.
+ ///
+ /// If there is a registered waker, it will be removed from the header and woken.
+ #[inline]
+ pub(crate) fn notify(&self) {
+ if let Some(waker) = self.swap_awaiter(None) {
+ // We need a safeguard against panics because waking can panic.
+ abort_on_panic(|| {
+ waker.wake();
+ });
+ }
+ }
+
+ /// Notifies the task blocked on the task unless its waker matches `current`.
+ ///
+ /// If there is a registered waker, it will be removed from the header.
+ #[inline]
+ pub(crate) fn notify_unless(&self, current: &Waker) {
+ if let Some(waker) = self.swap_awaiter(None) {
+ if !waker.will_wake(current) {
+ // We need a safeguard against panics because waking can panic.
+ abort_on_panic(|| {
+ waker.wake();
+ });
+ }
+ }
+ }
+
+ /// Swaps the awaiter and returns the previous value.
+ #[inline]
+ pub(crate) fn swap_awaiter(&self, new: Option<Waker>) -> Option<Waker> {
+ let new_is_none = new.is_none();
+
+ // We're about to try acquiring the lock in a loop. If it's already being held by another
+ // thread, we'll have to spin for a while so it's best to employ a backoff strategy.
+ let backoff = Backoff::new();
+ loop {
+ // Acquire the lock. If we're storing an awaiter, then also set the awaiter flag.
+ let state = if new_is_none {
+ self.state.fetch_or(LOCKED, Ordering::Acquire)
+ } else {
+ self.state.fetch_or(LOCKED | AWAITER, Ordering::Acquire)
+ };
+
+ // If the lock was acquired, break from the loop.
+ if state & LOCKED == 0 {
+ break;
+ }
+
+ // Snooze for a little while because the lock is held by another thread.
+ backoff.snooze();
+ }
+
+ // Replace the awaiter.
+ let old = self.awaiter.replace(new);
+
+ // Release the lock. If we've cleared the awaiter, then also unset the awaiter flag.
+ if new_is_none {
+ self.state.fetch_and(!LOCKED & !AWAITER, Ordering::Release);
+ } else {
+ self.state.fetch_and(!LOCKED, Ordering::Release);
+ }
+
+ old
+ }
+
+ /// Returns the offset at which the tag of type `T` is stored.
+ #[inline]
+ pub(crate) fn offset_tag<T>() -> usize {
+ let layout_header = Layout::new::<Header>();
+ let layout_t = Layout::new::<T>();
+ let (_, offset_t) = extend(layout_header, layout_t);
+ offset_t
+ }
+}
+
+impl fmt::Debug for Header {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ let state = self.state.load(Ordering::SeqCst);
+
+ f.debug_struct("Header")
+ .field("scheduled", &(state & SCHEDULED != 0))
+ .field("running", &(state & RUNNING != 0))
+ .field("completed", &(state & COMPLETED != 0))
+ .field("closed", &(state & CLOSED != 0))
+ .field("awaiter", &(state & AWAITER != 0))
+ .field("handle", &(state & HANDLE != 0))
+ .field("ref_count", &(state / REFERENCE))
+ .finish()
+ }
+}