ntex_rt/

task.rs

use std::{future::Future, pin::Pin, sync::Arc, sync::atomic::AtomicUsize, sync::atomic::Ordering};
use std::{task::Context, task::Poll};

// 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);
    }
}

/// Spawns a local future and returns a [`SendableJoinHandle`] for its result.
///
/// The task is spawned immediately on the current runtime, using
/// [`crate::spawn_detached`] and its registered task callbacks. The spawned
/// future need not be `Send`, but its output must be. Only the returned handle
/// may be moved to another thread; the task stays on its original runtime,
/// which must keep being driven.
///
/// Dropping the handle detaches the task unless
/// [`cancel_on_drop`](SendableJoinHandle::cancel_on_drop) is enabled.
/// Use [`cancel`](SendableJoinHandle::cancel) to request cancellation explicitly.
///
/// Awaiting the handle yields the output or [`crate::JoinError`] if
/// the task ends without one, for example on panic or runtime shutdown. Task
/// panics follow the runtime's usual behavior.
///
/// # Panics
///
/// Panics outside a runtime context that supports [`crate::spawn_detached`].
#[track_caller]
pub fn spawn_sendable<F>(fut: F) -> SendableJoinHandle<F::Output>
where
    F: Future + 'static,
    F::Output: Send,
{
    let (tx, rx) = oneshot::async_channel();
    let (cancel, canceled) = oneshot::async_channel::<()>();
    crate::spawn_detached(async move {
        let output = {
            let mut fut = std::pin::pin!(fut);
            let mut canceled = Some(canceled);
            std::future::poll_fn(|cx| {
                if let Some(rx) = &mut canceled {
                    match Pin::new(rx).poll(cx) {
                        Poll::Ready(Ok(())) => return Poll::Ready(None),
                        Poll::Ready(Err(_)) => canceled = None,
                        Poll::Pending => {}
                    }
                }
                fut.as_mut().poll(cx).map(Some)
            })
            .await
        };
        if let Some(output) = output {
            let _ = tx.send(output);
        }
    });
    SendableJoinHandle {
        rx,
        cancel: Some(cancel),
        cancel_on_drop: false,
    }
}

/// A sendable task handle returned by [`spawn_sendable()`].
///
/// The handle is `Send` when `T` is `Send`, but is not `Sync`. It may be moved
/// and polled on another thread while the task stays on its original runtime.
/// That runtime must keep being driven for the task to make progress.
///
/// Dropping the handle detaches the task unless
/// [`cancel_on_drop`](Self::cancel_on_drop) is enabled. Awaiting it yields the
/// output, or [`crate::JoinError`] if the task ends without one, including on
/// panic or runtime shutdown.
#[derive(Debug)]
pub struct SendableJoinHandle<T> {
    rx: oneshot::AsyncReceiver<T>,
    cancel: Option<oneshot::Sender<()>>,
    cancel_on_drop: bool,
}

impl<T> SendableJoinHandle<T> {
    /// Requests cancellation when this handle is dropped instead of detaching the task.
    ///
    /// Cancellation follows [`cancel`](Self::cancel) and does not wait for cleanup.
    /// A running poll is not interrupted.
    #[must_use]
    pub fn cancel_on_drop(mut self) -> Self {
        self.cancel_on_drop = true;
        self
    }

    /// Requests cancellation and discards the result without waiting for cleanup.
    ///
    /// The future is dropped on its runtime thread when the runtime next
    /// processes the task, or during shutdown. A running poll is not interrupted.
    /// This method may be called from another thread without a runtime context.
    pub fn cancel(mut self) {
        self.cancel_on_drop = true;
    }

    /// Returns `true` if the task completed or was closed without an output.
    pub fn is_finished(&self) -> bool {
        self.rx.has_message() || self.rx.is_closed()
    }
}

impl<T> Drop for SendableJoinHandle<T> {
    fn drop(&mut self) {
        if self.cancel_on_drop
            && let Some(cancel) = self.cancel.take()
        {
            let _ = cancel.send(());
        }
    }
}

impl<T> Future for SendableJoinHandle<T> {
    type Output = Result<T, crate::JoinError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        Pin::new(&mut self.rx)
            .poll(cx)
            .map(|result| result.map_err(|_| crate::JoinError))
    }
}

/// Installs process-wide callbacks for instrumenting spawned tasks.
///
/// Callbacks apply when [`crate::spawn`], [`crate::spawn_detached`],
/// [`spawn_sendable`], [`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;