some_executor 0.7.2

// SPDX-License-Identifier: MIT OR Apache-2.0

//! The per-task observation state, and what it costs when nobody is watching.
//!
//! A spawned task owns one [`TaskInstrumentation`]: its causal context, its
//! ID, its cancellation handle, and — behind the `logwise-forensic` and
//! `logwise-performance` features — the wake bookkeeping needed to report wake
//! latency and active poll time as separate measurements.
//!
//! Nearly every field here is `#[cfg]`-gated on one of those features. That is
//! the point: a build with the instrumentation off does not pay for an
//! `Arc<WakeState>`, a `created_at` timestamp, or the atomics behind them,
//! because the fields are not compiled at all rather than present and unused.

use super::{InFlightTaskCancellation, TaskID};
use logwise::ContextToken;
#[cfg(feature = "logwise-performance")]
use std::time::Duration;
#[cfg(any(feature = "logwise-forensic", feature = "logwise-performance"))]
use std::{
    sync::{Arc, Mutex},
    task::{Wake, Waker},
};

/// Durable causal and observation state owned by a spawned task.
pub(super) struct TaskInstrumentation {
    context: ContextToken,
    task_id: TaskID,
    cancellation: InFlightTaskCancellation,
    first_poll: bool,
    terminal: Option<Terminal>,
    #[cfg(feature = "logwise-performance")]
    created_at: crate::sys::Instant,
    #[cfg(any(feature = "logwise-forensic", feature = "logwise-performance"))]
    wake: Arc<WakeState>,
}

#[derive(Clone, Copy)]
enum Terminal {
    Completed,
    Cancelled,
}

impl Terminal {
    const fn name(self) -> &'static str {
        match self {
            Self::Completed => "completed",
            Self::Cancelled => "cancelled",
        }
    }
}

impl TaskInstrumentation {
    pub(super) fn new(
        task_id: TaskID,
        label: &str,
        #[cfg_attr(not(feature = "exfiltrate"), allow(unused_variables))]
        priority: priority::Priority,
        #[cfg_attr(not(feature = "exfiltrate"), allow(unused_variables))] hint: crate::hint::Hint,
        cancellation: InFlightTaskCancellation,
    ) -> Self {
        #[cfg(feature = "exfiltrate")]
        super::registry::record_spawn(task_id, label, priority, hint);
        let parent = logwise::context::capture();
        let context = logwise::context::child(parent, "some_executor.task");
        let instrumentation = Self {
            context,
            task_id,
            cancellation,
            first_poll: true,
            terminal: None,
            #[cfg(feature = "logwise-performance")]
            created_at: crate::sys::Instant::now(),
            #[cfg(any(feature = "logwise-forensic", feature = "logwise-performance"))]
            wake: Arc::new(WakeState::new(context, task_id)),
        };
        instrumentation.spawned(label);
        instrumentation
    }

    pub(super) fn context(&self) -> ContextToken {
        self.context
    }

    pub(super) fn begin_poll(&mut self) -> ActivePoll {
        #[cfg(feature = "exfiltrate")]
        super::registry::record_poll(self.task_id);
        if self.first_poll {
            self.first_poll = false;
            let _entered = logwise::context::enter(self.context);
            #[cfg(feature = "logwise-forensic")]
            logwise::forensic!(
                "some_executor.task.first_poll",
                task_id = support(self.task_id.to_u64()),
                poll_number = support(1_u64),
            );
        }

        #[cfg(feature = "logwise-performance")]
        if let Some(woken_at) = self.wake.take_wake() {
            let _entered = logwise::context::enter(self.context);
            logwise::measurement!(
                "some_executor.task.wake_latency",
                task_id = support(self.task_id.to_u64()),
                duration_ns = support(duration_ns(woken_at.elapsed())),
            );
        }

        ActivePoll {
            context: self.context,
            task_id: self.task_id,
            #[cfg(feature = "logwise-performance")]
            started_at: crate::sys::Instant::now(),
        }
    }

    #[cfg(any(feature = "logwise-forensic", feature = "logwise-performance"))]
    pub(super) fn instrument_waker(&self, waker: &Waker) -> Waker {
        Waker::from(Arc::new(ObservedWaker {
            inner: waker.clone(),
            observation: self.wake.clone(),
        }))
    }

    pub(super) fn cancelled(&mut self) {
        if self.terminal.is_some() {
            return;
        }
        self.terminal = Some(Terminal::Cancelled);
        #[cfg(feature = "exfiltrate")]
        super::registry::record_terminal(self.task_id, super::registry::TaskState::Cancelled);
        let _entered = logwise::context::enter(self.context);
        #[cfg(feature = "logwise-forensic")]
        logwise::forensic!(
            "some_executor.task.cancelled",
            task_id = support(self.task_id.to_u64()),
        );
        self.wall_measurement(Terminal::Cancelled.name());
    }

    pub(super) fn completed(&mut self) {
        if self.cancellation.is_cancelled() {
            self.cancelled();
            return;
        }
        if self.terminal.is_some() {
            return;
        }
        self.terminal = Some(Terminal::Completed);
        #[cfg(feature = "exfiltrate")]
        super::registry::record_terminal(self.task_id, super::registry::TaskState::Completed);
        let _entered = logwise::context::enter(self.context);
        #[cfg(feature = "logwise-forensic")]
        logwise::forensic!(
            "some_executor.task.completed",
            task_id = support(self.task_id.to_u64()),
        );
        self.wall_measurement(Terminal::Completed.name());
    }

    fn spawned(&self, label: &str) {
        let _entered = logwise::context::enter(self.context);
        #[cfg(all(feature = "logwise-forensic", feature = "logwise-diagnostic"))]
        logwise::forensic!(
            "some_executor.task.spawned",
            task_id = support(self.task_id.to_u64()),
            detail task_label = local(label),
        );
        #[cfg(all(feature = "logwise-forensic", not(feature = "logwise-diagnostic")))]
        logwise::forensic!(
            "some_executor.task.spawned",
            task_id = support(self.task_id.to_u64()),
        );
        #[cfg(not(feature = "logwise-forensic"))]
        let _ = label;
    }

    fn wall_measurement(&self, state: &'static str) {
        #[cfg(feature = "logwise-performance")]
        logwise::measurement!(
            "some_executor.task.wall_lifetime",
            task_id = support(self.task_id.to_u64()),
            state = support(state),
            duration_ns = support(duration_ns(self.created_at.elapsed())),
        );
        #[cfg(not(feature = "logwise-performance"))]
        let _ = state;
    }
}

impl Drop for TaskInstrumentation {
    fn drop(&mut self) {
        if self.terminal.is_none() && self.cancellation.is_cancelled() {
            self.cancelled();
        }
        let state = self.terminal.map_or("dropped", Terminal::name);
        let _entered = logwise::context::enter(self.context);
        #[cfg(feature = "logwise-forensic")]
        logwise::forensic!(
            "some_executor.task.dropped",
            task_id = support(self.task_id.to_u64()),
            state = support(state),
        );
        if self.terminal.is_none() {
            self.wall_measurement(state);
        }
    }
}

pub(super) struct ActivePoll {
    context: ContextToken,
    task_id: TaskID,
    #[cfg(feature = "logwise-performance")]
    started_at: crate::sys::Instant,
}

impl Drop for ActivePoll {
    fn drop(&mut self) {
        #[cfg(feature = "logwise-performance")]
        {
            let _entered = logwise::context::enter(self.context);
            logwise::measurement!(
                "some_executor.task.active_poll_time",
                task_id = support(self.task_id.to_u64()),
                duration_ns = support(duration_ns(self.started_at.elapsed())),
            );
        }
        #[cfg(not(feature = "logwise-performance"))]
        let _ = (self.context, self.task_id);
    }
}

#[cfg(any(feature = "logwise-forensic", feature = "logwise-performance"))]
struct WakeState {
    context: ContextToken,
    task_id: TaskID,
    #[cfg(feature = "logwise-performance")]
    woken_at: Mutex<Option<crate::sys::Instant>>,
}

#[cfg(any(feature = "logwise-forensic", feature = "logwise-performance"))]
impl WakeState {
    fn new(context: ContextToken, task_id: TaskID) -> Self {
        Self {
            context,
            task_id,
            #[cfg(feature = "logwise-performance")]
            woken_at: Mutex::new(None),
        }
    }

    fn observed(&self) {
        #[cfg(feature = "logwise-performance")]
        {
            let mut woken_at = self
                .woken_at
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            if woken_at.is_none() {
                *woken_at = Some(crate::sys::Instant::now());
            }
        }
        let _entered = logwise::context::enter(self.context);
        #[cfg(feature = "logwise-forensic")]
        logwise::forensic!(
            "some_executor.task.woken",
            task_id = support(self.task_id.to_u64()),
        );
    }

    #[cfg(feature = "logwise-performance")]
    fn take_wake(&self) -> Option<crate::sys::Instant> {
        self.woken_at
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .take()
    }
}

#[cfg(any(feature = "logwise-forensic", feature = "logwise-performance"))]
struct ObservedWaker {
    inner: Waker,
    observation: Arc<WakeState>,
}

#[cfg(any(feature = "logwise-forensic", feature = "logwise-performance"))]
impl Wake for ObservedWaker {
    fn wake(self: Arc<Self>) {
        self.observation.observed();
        self.inner.wake_by_ref();
    }

    fn wake_by_ref(self: &Arc<Self>) {
        self.observation.observed();
        self.inner.wake_by_ref();
    }
}

#[cfg(feature = "logwise-performance")]
fn duration_ns(duration: Duration) -> u64 {
    duration.as_nanos().min(u128::from(u64::MAX)) as u64
}