ax-task 0.8.2

OS-independent IRQ-safe SMP task scheduling core
Documentation
//! Allocation-free raw waker implementation.

use core::{
    mem::ManuallyDrop,
    ptr,
    task::{RawWaker, RawWakerVTable, Waker},
};

use super::coroutine::{
    CoroutineHeader, release_reference, retain_reference, schedule, schedule_sync,
};

static VTABLE: RawWakerVTable = RawWakerVTable::new(clone_waker, wake, wake_by_ref, drop_waker);

/// Creates an owning standard waker for one pinned coroutine header.
///
/// # Safety
///
/// `header` must be pinned and live, and the caller must own a reference that
/// keeps it valid until this function retains the waker's independent reference.
pub(super) unsafe fn coroutine_waker(header: *mut CoroutineHeader) -> Waker {
    let header_ref = unsafe {
        // Caller owns a live reference and the pinned header remains valid while
        // the returned Waker owns the reference retained below.
        &*header
    };
    retain_reference(header_ref);
    let raw = RawWaker::new(header.cast(), &VTABLE);
    unsafe {
        // `raw` owns exactly the reference retained above and uses the matching
        // vtable to release it.
        Waker::from_raw(raw)
    }
}

/// Consumes a Waker with Linux `WF_SYNC` semantics when it belongs to this
/// executor, otherwise preserves the standard Waker contract.
pub(super) fn wake_sync(waker: Waker) {
    if !ptr::eq(waker.vtable(), &VTABLE) {
        waker.wake();
        return;
    }

    let header = waker.data().cast_mut().cast::<CoroutineHeader>();
    let _waker = ManuallyDrop::new(waker);
    unsafe {
        // Matching this private vtable proves that the data pointer owns one
        // coroutine-header reference. This path consumes and releases exactly
        // the same reference as the ordinary raw `wake` callback.
        schedule_sync(header);
        release_reference(header);
    }
}

/// Clones one raw-waker reference.
///
/// # Safety
///
/// `data` must originate from this module's vtable and own a live header reference.
unsafe fn clone_waker(data: *const ()) -> RawWaker {
    let header = data.cast_mut().cast::<CoroutineHeader>();
    let header_ref = unsafe {
        // Every call is made through a live RawWaker that owns one reference.
        &*header
    };
    retain_reference(header_ref);
    RawWaker::new(data, &VTABLE)
}

/// Publishes a wake and consumes one raw-waker reference.
///
/// # Safety
///
/// `data` must originate from this module's vtable and own a live header reference.
unsafe fn wake(data: *const ()) {
    let header = data.cast_mut().cast::<CoroutineHeader>();
    unsafe {
        // The consumed RawWaker keeps the header live through publication. The
        // newly queued node takes its own reference before this one is released.
        schedule(header);
        release_reference(header);
    }
}

/// Publishes a wake while retaining the borrowed raw-waker reference.
///
/// # Safety
///
/// `data` must originate from this module's vtable and own a live header reference.
unsafe fn wake_by_ref(data: *const ()) {
    let header = data.cast_mut().cast::<CoroutineHeader>();
    unsafe {
        // The borrowed RawWaker retains its reference after publication returns.
        schedule(header);
    }
}

/// Releases one raw-waker reference.
///
/// # Safety
///
/// `data` must originate from this module's vtable and own a live header reference.
unsafe fn drop_waker(data: *const ()) {
    let header = data.cast_mut().cast::<CoroutineHeader>();
    unsafe {
        // Task context reclaims the final reference immediately. Hard IRQ only
        // publishes the embedded typed reclaim node and never runs destructors.
        release_reference(header);
    }
}