ntex_rt/

task.rs

use std::{future::Future, sync::Arc, sync::atomic::AtomicUsize, sync::atomic::Ordering};

// The Callbacks static holds a pointer to the global callbacks. It is protected by
// the STATE static which determines whether `CBS` has been initialized yet.
static mut CBS: Option<Arc<dyn CallbacksApi>> = None;

static STATE: AtomicUsize = AtomicUsize::new(0);

// There are three different states that we care about: the callback's
// uninitialized, the callback's initializing (set_cbs's been called but
// CBS hasn't actually been set yet), or the callbacks's active.
const UNINITIALIZED: usize = 0;
const INITIALIZING: usize = 1;
const INITIALIZED: usize = 2;

trait CallbacksApi {
    fn before(&self) -> Option<*const ()>;
    fn enter(&self, _: *const ()) -> *const ();
    fn exit(&self, _: *const ());
    fn after(&self, _: *const ());
}

#[allow(clippy::struct_field_names)]
struct Callbacks<A, B, C, D> {
    f_before: A,
    f_enter: B,
    f_exit: C,
    f_after: D,
}

impl<A, B, C, D> CallbacksApi for Callbacks<A, B, C, D>
where
    A: Fn() -> Option<*const ()> + 'static,
    B: Fn(*const ()) -> *const () + 'static,
    C: Fn(*const ()) + 'static,
    D: Fn(*const ()) + 'static,
{
    fn before(&self) -> Option<*const ()> {
        (self.f_before)()
    }
    fn enter(&self, d: *const ()) -> *const () {
        (self.f_enter)(d)
    }
    fn exit(&self, d: *const ()) {
        (self.f_exit)(d);
    }
    fn after(&self, d: *const ()) {
        (self.f_after)(d);
    }
}

pub(crate) struct Data {
    cb: &'static dyn CallbacksApi,
    ptr: *const (),
}

impl Data {
    #[allow(clippy::if_not_else)]
    pub(crate) fn load() -> Option<Data> {
        // `Acquire` pairs with the `Release` store in `set_cbs`, so `CBS` is visible
        let cb = if STATE.load(Ordering::Acquire) != INITIALIZED {
            None
        } else {
            #[allow(static_mut_refs)]
            unsafe {
                Some(CBS.as_ref().map(AsRef::as_ref).unwrap())
            }
        };

        if let Some(cb) = cb
            && let Some(ptr) = cb.before()
        {
            return Some(Data { cb, ptr });
        }
        None
    }

    pub(crate) fn run<F, R>(&mut self, f: F) -> R
    where
        F: FnOnce() -> R,
    {
        struct Exit(&'static dyn CallbacksApi, *const ());

        impl Drop for Exit {
            fn drop(&mut self) {
                self.0.exit(self.1);
            }
        }

        // `exit` also runs if `f` panics, so the context set by `enter` is restored
        let _exit = Exit(self.cb, self.cb.enter(self.ptr));
        f()
    }
}

/// Wraps a future so that task callbacks, if registered, run around each poll.
///
/// Always returns the same future type, so spawning code is generated once
/// per future rather than once per callback mode.
pub(crate) fn wrap<F: Future>(fut: F) -> impl Future<Output = F::Output> {
    let mut data = Data::load();
    async move {
        let mut f = std::pin::pin!(fut);
        std::future::poll_fn(|cx| match data.as_mut() {
            Some(data) => data.run(|| f.as_mut().poll(cx)),
            None => f.as_mut().poll(cx),
        })
        .await
    }
}

impl Drop for Data {
    fn drop(&mut self) {
        self.cb.after(self.ptr);
    }
}

/// Installs process-wide callbacks for instrumenting spawned tasks.
///
/// Callbacks apply when [`crate::spawn`], [`crate::spawn_detached`],
/// [`crate::rt::spawn`], or [`crate::rt::spawn_detached`] wraps a future, not
/// retroactively to existing tasks. Direct [`crate::rt::Runtime`] spawning
/// methods do not apply them.
///
/// - `f_before` runs synchronously when the future is wrapped. Returning `None`
///   disables the remaining callbacks for that future.
/// - `f_enter` receives the task's pointer before each poll and returns a pointer
///   representing the previous context.
/// - `f_exit` receives that previous-context pointer after the poll, including
///   when polling unwinds.
/// - `f_after` receives the task's original pointer when its callback state is
///   dropped, including cancellation before the first poll.
///
/// Registration is permanent and first-wins. Does nothing if callbacks were
/// already installed; use [`task_opt_callbacks`] to observe registration success.
///
/// # Safety
///
/// Pointers are opaque to the runtime. The caller must keep the state returned
/// by `f_before` valid until `f_after`, and the state returned by `f_enter` valid
/// until its matching `f_exit`. Callbacks must safely handle concurrent use by
/// different runtime threads and any nested task activity they trigger.
pub unsafe fn task_callbacks<FBefore, FEnter, FExit, FAfter>(
    f_before: FBefore,
    f_enter: FEnter,
    f_exit: FExit,
    f_after: FAfter,
) where
    FBefore: Fn() -> Option<*const ()> + 'static + Sync,
    FEnter: Fn(*const ()) -> *const () + 'static + Sync,
    FExit: Fn(*const ()) + 'static + Sync,
    FAfter: Fn(*const ()) + 'static + Sync,
{
    let new = Arc::new(Callbacks {
        f_before,
        f_enter,
        f_exit,
        f_after,
    });
    let _ = set_cbs(new);
}

/// Attempts to install process-wide task callbacks.
///
/// Uses the callback lifecycle described by [`task_callbacks`]. Returns `true`
/// if this call installs the callbacks, or `false` if another registration has
/// already won. Installed callbacks cannot be replaced or removed.
///
/// # Safety
///
/// The pointer lifetimes and callback thread-safety requirements of
/// [`task_callbacks`] apply.
pub unsafe fn task_opt_callbacks<FBefore, FEnter, FExit, FAfter>(
    f_before: FBefore,
    f_enter: FEnter,
    f_exit: FExit,
    f_after: FAfter,
) -> bool
where
    FBefore: Fn() -> Option<*const ()> + Sync + 'static,
    FEnter: Fn(*const ()) -> *const () + Sync + 'static,
    FExit: Fn(*const ()) + Sync + 'static,
    FAfter: Fn(*const ()) + Sync + 'static,
{
    let new = Arc::new(Callbacks {
        f_before,
        f_enter,
        f_exit,
        f_after,
    });
    set_cbs(new).is_ok()
}

fn set_cbs(cbs: Arc<dyn CallbacksApi>) -> Result<(), ()> {
    match STATE.compare_exchange(
        UNINITIALIZED,
        INITIALIZING,
        Ordering::Acquire,
        Ordering::Relaxed,
    ) {
        Ok(UNINITIALIZED) => {
            unsafe {
                CBS = Some(cbs);
            }
            STATE.store(INITIALIZED, Ordering::Release);
            Ok(())
        }
        Err(INITIALIZING) => {
            while STATE.load(Ordering::Relaxed) == INITIALIZING {
                std::hint::spin_loop();
            }
            Err(())
        }
        _ => Err(()),
    }
}

#[cfg(test)]
#[path = "../tests/unit/task.rs"]
mod tests;