some_executor 0.7.2

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

//! A bounded record of the tasks this process has spawned.
//!
//! "Which tasks exist and what are they doing" is the first question when an
//! async program hangs, and until now the answer needed a debugger. This keeps
//! enough to answer it, and no more.
//!
//! # Bounded, and honest about it
//!
//! The registry keeps the most recent `N` tasks and nothing else. `N` defaults
//! to [`DEFAULT_CAPACITY`] and can be set with `SOME_EXECUTOR_TASK_REGISTRY_CAPACITY`.
//!
//! When it overflows it drops the **oldest terminal** entry first, and only
//! falls back to dropping the oldest live one when every entry is live. That
//! ordering matters: the oldest un-completed task is usually the bug, so it is
//! the last thing worth forgetting.
//!
//! Whatever it drops, it counts. A snapshot that quietly returned fewer tasks
//! than exist would be worse than no snapshot, because a caller cannot tell
//! "nothing else is running" from "I lost it" — so the count is reported and
//! the answer is marked partial.
//!
//! # Cost
//!
//! Behind the `exfiltrate` feature. With it off, none of this is compiled and
//! spawning does not touch a lock.

use std::collections::VecDeque;
use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering};

use crate::hint::Hint;
use crate::sys::Instant;
use crate::task::TaskID;
use priority::Priority;

/// Kept when `SOME_EXECUTOR_TASK_REGISTRY_CAPACITY` is unset or unparseable.
pub const DEFAULT_CAPACITY: usize = 1024;

/// What a task is doing now.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum TaskState {
    /// Spawned, not finished. Never polled yet if `polls` is zero.
    Live,
    /// Ran to completion and produced its value.
    Completed,
    /// Dropped or cancelled before producing one.
    Cancelled,
}

impl TaskState {
    /// The stable wire name reported by `snapshot`, matched by external
    /// tooling rather than rendered from `Debug`.
    pub const fn name(self) -> &'static str {
        match self {
            TaskState::Live => "live",
            TaskState::Completed => "completed",
            TaskState::Cancelled => "cancelled",
        }
    }
}

/// One spawned task, as an observer outside the process sees it.
#[derive(Clone, Debug)]
pub struct TaskEntry {
    /// The task's own id. What `--id` selects.
    pub task_id: TaskID,
    /// The label the spawner gave it. Caller-derived, so it is local-only.
    pub label: String,
    /// The priority it was spawned at.
    pub priority: Priority,
    /// The scheduling hint it was spawned with.
    pub hint: Hint,
    /// When it was spawned. Against `last_poll` this is what identifies a task
    /// that was spawned and then never got a turn.
    pub spawned_at: Instant,
    /// What it is doing now.
    pub state: TaskState,
    /// How many times it has been polled. Zero on a `Live` task means nothing
    /// has even started it -- a different bug from one whose wake never came.
    pub polls: u64,
    /// When it was last polled, and `None` if it never has been.
    pub last_poll: Option<Instant>,
}

struct Registry {
    capacity: usize,
    entries: VecDeque<TaskEntry>,
}

static REGISTRY: Mutex<Option<Registry>> = Mutex::new(None);
/// Counted rather than inferred: a caller must be able to tell "nothing else is
/// running" from "I lost it".
static DROPPED: AtomicU64 = AtomicU64::new(0);
static SPAWNED: AtomicU64 = AtomicU64::new(0);
static COMPLETED: AtomicU64 = AtomicU64::new(0);
static CANCELLED: AtomicU64 = AtomicU64::new(0);

fn configured_capacity() -> usize {
    std::env::var("SOME_EXECUTOR_TASK_REGISTRY_CAPACITY")
        .ok()
        .and_then(|value| value.parse::<usize>().ok())
        .filter(|capacity| *capacity > 0)
        .unwrap_or(DEFAULT_CAPACITY)
}

fn with<R>(f: impl FnOnce(&mut Registry) -> R) -> Option<R> {
    // `try_lock`, deliberately: a spawn must not block on a query, and a query
    // must not block a spawn. A missed record is better than a stalled
    // executor, and the miss is counted.
    let mut guard = REGISTRY.try_lock().ok()?;
    let registry = guard.get_or_insert_with(|| Registry {
        capacity: configured_capacity(),
        entries: VecDeque::new(),
    });
    Some(f(registry))
}

/// Records a spawn. Called on the spawning thread, so it must be cheap.
pub(crate) fn record_spawn(task_id: TaskID, label: &str, priority: Priority, hint: Hint) {
    SPAWNED.fetch_add(1, Ordering::Relaxed);
    let recorded = with(|registry| {
        while registry.entries.len() >= registry.capacity {
            // Prefer forgetting something that already finished.
            let victim = registry
                .entries
                .iter()
                .position(|entry| entry.state != TaskState::Live)
                .unwrap_or(0);
            registry.entries.remove(victim);
            DROPPED.fetch_add(1, Ordering::Relaxed);
        }
        registry.entries.push_back(TaskEntry {
            task_id,
            label: label.to_string(),
            priority,
            hint,
            spawned_at: Instant::now(),
            state: TaskState::Live,
            polls: 0,
            last_poll: None,
        });
    });
    if recorded.is_none() {
        DROPPED.fetch_add(1, Ordering::Relaxed);
    }
}

pub(crate) fn record_poll(task_id: TaskID) {
    let _ = with(|registry| {
        if let Some(entry) = registry
            .entries
            .iter_mut()
            .find(|entry| entry.task_id == task_id)
        {
            entry.polls += 1;
            entry.last_poll = Some(Instant::now());
        }
    });
}

pub(crate) fn record_terminal(task_id: TaskID, state: TaskState) {
    match state {
        TaskState::Completed => COMPLETED.fetch_add(1, Ordering::Relaxed),
        TaskState::Cancelled => CANCELLED.fetch_add(1, Ordering::Relaxed),
        TaskState::Live => 0,
    };
    let _ = with(|registry| {
        if let Some(entry) = registry
            .entries
            .iter_mut()
            .find(|entry| entry.task_id == task_id)
        {
            entry.state = state;
        }
    });
}

/// Lifetime counters. Unlike the entries, these are never dropped.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct RegistryStats {
    /// Tasks spawned over the process's life.
    pub spawned: u64,
    /// How many of those ran to completion.
    pub completed: u64,
    /// How many were cancelled or dropped first.
    pub cancelled: u64,
    /// Entries the registry forgot, either to stay bounded or because it could
    /// not take its lock.
    pub dropped: u64,
    /// How many entries are held right now.
    pub retained: usize,
    /// The bound currently in force.
    pub capacity: usize,
}

/// Reads the lifetime counters without copying out the entries.
pub fn stats() -> RegistryStats {
    let (retained, capacity) = with(|registry| (registry.entries.len(), registry.capacity))
        .unwrap_or((0, configured_capacity()));
    RegistryStats {
        spawned: SPAWNED.load(Ordering::Relaxed),
        completed: COMPLETED.load(Ordering::Relaxed),
        cancelled: CANCELLED.load(Ordering::Relaxed),
        dropped: DROPPED.load(Ordering::Relaxed),
        retained,
        capacity,
    }
}

/// The retained entries, oldest first — the order a reader wants, because the
/// oldest un-completed task is usually the bug.
///
/// `None` when the registry was locked; the caller should report `busy` rather
/// than an empty list.
pub fn entries() -> Option<Vec<TaskEntry>> {
    with(|registry| registry.entries.iter().cloned().collect())
}

/// One lock for every test that touches the registry, wherever it lives.
///
/// The registry is process-global, so a per-module lock is not enough: two
/// modules with their own locks interleave and stomp each other, which is
/// exactly what happened first time.
#[cfg(test)]
pub(crate) static TEST_SERIALIZE: Mutex<()> = Mutex::new(());

/// Takes the shared lock and resets the registry.
#[cfg(test)]
pub(crate) fn test_session(capacity: usize) -> std::sync::MutexGuard<'static, ()> {
    let guard = TEST_SERIALIZE
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner());
    reset_for_test(capacity);
    guard
}

#[cfg(test)]
pub(crate) fn reset_for_test(capacity: usize) {
    let mut guard = REGISTRY.lock().unwrap_or_else(|e| e.into_inner());
    *guard = Some(Registry {
        capacity,
        entries: VecDeque::new(),
    });
    DROPPED.store(0, Ordering::Relaxed);
    SPAWNED.store(0, Ordering::Relaxed);
    COMPLETED.store(0, Ordering::Relaxed);
    CANCELLED.store(0, Ordering::Relaxed);
}

#[cfg(test)]
mod tests {
    use super::*;

    use super::test_session as session;

    fn spawn(id: u64, label: &str) {
        record_spawn(
            TaskID::from_u64(id),
            label,
            Priority::unit_test(),
            Hint::Unknown,
        );
    }

    #[test]
    fn a_spawn_is_recorded_and_starts_live_and_unpolled() {
        let _session = session(8);
        spawn(1, "first");
        let entries = entries().expect("registry available");
        assert_eq!(entries.len(), 1);
        assert_eq!(entries[0].state, TaskState::Live);
        assert_eq!(entries[0].polls, 0);
        assert!(
            entries[0].last_poll.is_none(),
            "never polled is distinguishable from polled long ago"
        );
        assert_eq!(stats().spawned, 1);
    }

    #[test]
    fn polls_and_terminal_states_are_tracked() {
        let _session = session(8);
        spawn(1, "first");
        record_poll(TaskID::from_u64(1));
        record_poll(TaskID::from_u64(1));
        record_terminal(TaskID::from_u64(1), TaskState::Completed);

        let entries = entries().unwrap();
        assert_eq!(entries[0].polls, 2);
        assert!(entries[0].last_poll.is_some());
        assert_eq!(entries[0].state, TaskState::Completed);
        assert_eq!(stats().completed, 1);
        assert_eq!(stats().cancelled, 0);
    }

    /// Overflow drops a *finished* task before a live one: the oldest
    /// un-completed task is usually the bug, so it is the last thing worth
    /// forgetting.
    #[test]
    fn overflow_forgets_finished_tasks_before_live_ones() {
        let _session = session(2);
        spawn(1, "live-and-old");
        spawn(2, "finished");
        record_terminal(TaskID::from_u64(2), TaskState::Completed);

        spawn(3, "newcomer");

        let entries = entries().unwrap();
        let ids: Vec<u64> = entries.iter().map(|e| e.task_id.to_u64()).collect();
        assert!(
            ids.contains(&1),
            "the live task should have survived: {ids:?}"
        );
        assert!(ids.contains(&3));
        assert!(!ids.contains(&2), "the finished one should have gone first");
    }

    /// When everything is live it still has to drop something — the oldest —
    /// and it must say so.
    #[test]
    fn overflow_of_all_live_tasks_drops_the_oldest_and_counts_it() {
        let _session = session(2);
        spawn(1, "oldest");
        spawn(2, "middle");
        spawn(3, "newest");

        let entries = entries().unwrap();
        let ids: Vec<u64> = entries.iter().map(|e| e.task_id.to_u64()).collect();
        assert_eq!(ids, vec![2, 3]);
        assert_eq!(stats().dropped, 1, "a drop must be counted, not implied");
        assert_eq!(stats().retained, 2);
    }

    /// Lifetime counters survive what the bounded entries do not.
    #[test]
    fn counters_outlive_the_entries_they_counted() {
        let _session = session(1);
        for id in 1..=10 {
            spawn(id, "churn");
        }
        let stats = stats();
        assert_eq!(stats.spawned, 10);
        assert_eq!(stats.retained, 1);
        assert_eq!(stats.dropped, 9);
    }
}