harn-vm 0.10.131

//! One owner for running a Harn child interpreter as a subtask.
//!
//! `spawn`, `parallel`, `parallel each`, `parallel settle`, the parallel
//! stream fan-out, and `pool.submit` all create the same thing: a child
//! interpreter that runs a closure concurrently with its parent. Each one
//! needs the same three things, and each used to arrange them differently.
//!
//! 1. An isolated copy of the parent's ambient execution scope. It is captured
//!    eagerly, while the parent's scope is still swapped in, so the subtask
//!    carries the agent session, policies, and event attribution that were
//!    live at the moment it was created.
//! 2. The parent's pool registry, so `pool.*` inside the subtask reaches the
//!    same pools instead of a fresh per-thread fallback.
//! 3. A place on the runtime to run.
//!
//! [`prepare`] does the first two. [`spawn`] and [`spawn_into`] do the third.
//! Nothing else in the crate decides where a child interpreter runs.
//!
//! # Placement
//!
//! [`SubtaskPlacement`] selects between the creating thread and the runtime's
//! worker threads. Worker placement is the default now that every capability
//! reachable from a child interpreter has an execution-scoped cross-thread
//! owner. It gives CPU-bound fan-out real parallelism:
//! a tight compute loop never yields, so pinned subtasks run one at a time no
//! matter how many workers are idle.
//!
//! Both placements require the subtask future to be `Send + 'static`. That is
//! deliberate. The bound is a property of the seam, not of the placement, so
//! changing the default cannot turn a compiling program into a broken one.

use std::future::Future;
use std::sync::Arc;
use std::task::{Context, Poll};

use crate::orchestration::{scope_ambient, AmbientExecutionScope};
use crate::stdlib::pool::{with_pool_registry_scope, PoolRegistry};
use pin_project_lite::pin_project;

pin_project! {
    /// Proof that a future carries the ambient scope required at a runtime
    /// thread boundary.
    ///
    /// The constructor stays private to this module. Callers can only obtain
    /// one through [`prepare`], and the spawn functions only accept this type,
    /// so a new child-interpreter path cannot accidentally bypass scope
    /// capture while still compiling.
    pub(crate) struct PreparedSubtask<F> {
        #[pin]
        inner: F,
    }
}

impl<F: Future> Future for PreparedSubtask<F> {
    type Output = F::Output;

    fn poll(self: std::pin::Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
        self.project().inner.poll(context)
    }
}

impl<F: Future> PreparedSubtask<F> {
    /// Transform the output without discarding the proof carried by this
    /// future. Executors use this to attach source-order indices before
    /// inserting a branch into a `JoinSet`.
    pub(crate) fn map_output<M, T>(self, map: M) -> PreparedSubtask<impl Future<Output = T>>
    where
        M: FnOnce(F::Output) -> T,
    {
        PreparedSubtask {
            inner: async move { map(self.await) },
        }
    }
}

/// Where a subtask runs.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum SubtaskPlacement {
    /// Run on the thread that created the subtask. Branches interleave at
    /// await points and never migrate, preserving thread-affine capability and
    /// embedding-host contracts.
    CurrentThread,
    /// Run on the runtime's worker threads. CPU-bound branches of one fan-out
    /// then run at the same time on different cores. `current_thread` remains
    /// an explicit compatibility mode for deliberately single-threaded hosts.
    #[default]
    Worker,
}

/// The environment variable that selects placement for a whole process.
pub const PLACEMENT_ENV: &str = "HARN_VM_SUBTASK_PLACEMENT";
/// Canonical values accepted for [`PLACEMENT_ENV`]. The environment registry
/// consumes this same vocabulary, so startup validation and placement parsing
/// cannot drift.
pub const PLACEMENT_VALUES: &[&str] = &["worker", "current_thread"];

/// An operator supplied a placement outside the runtime-owned vocabulary.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SubtaskPlacementParseError {
    value: String,
}

impl std::fmt::Display for SubtaskPlacementParseError {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            formatter,
            "invalid {PLACEMENT_ENV} value {:?}; expected one of {}",
            self.value,
            PLACEMENT_VALUES.join(", ")
        )
    }
}

impl std::error::Error for SubtaskPlacementParseError {}

impl SubtaskPlacement {
    /// Parse an operator-supplied placement name against the runtime-owned
    /// closed vocabulary.
    pub fn from_env_value(value: &str) -> Result<Self, SubtaskPlacementParseError> {
        match value.trim().to_ascii_lowercase().as_str() {
            "worker" => Ok(Self::Worker),
            "current_thread" => Ok(Self::CurrentThread),
            _ => Err(SubtaskPlacementParseError {
                value: value.to_string(),
            }),
        }
    }

    fn name(self) -> &'static str {
        match self {
            Self::Worker => "worker",
            Self::CurrentThread => "current_thread",
        }
    }
}

impl std::fmt::Display for SubtaskPlacement {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter.write_str(self.name())
    }
}

/// The process-wide placement, read from the environment once.
fn placement_from_environment() -> SubtaskPlacement {
    static RESOLVED: std::sync::OnceLock<SubtaskPlacement> = std::sync::OnceLock::new();
    *RESOLVED.get_or_init(|| {
        let Ok(value) = std::env::var(PLACEMENT_ENV) else {
            return SubtaskPlacement::default();
        };
        SubtaskPlacement::from_env_value(&value).unwrap_or_else(|error| panic!("{error}"))
    })
}

thread_local! {
    /// An execution-scoped placement override. It is captured by
    /// [`AmbientExecutionScope`], so a nested subtask keeps the placement its
    /// execution tree was started with instead of falling back to the process
    /// default on a worker thread.
    static SUBTASK_PLACEMENT_CONTEXT: std::cell::RefCell<Option<SubtaskPlacement>> =
        const { std::cell::RefCell::new(None) };
}

/// Swap the execution-scoped placement override. Paired with
/// [`AmbientExecutionScope`]'s per-poll swap.
pub(crate) fn swap_subtask_placement_context(
    next: Option<SubtaskPlacement>,
) -> Option<SubtaskPlacement> {
    SUBTASK_PLACEMENT_CONTEXT.with(|slot| std::mem::replace(&mut *slot.borrow_mut(), next))
}

/// The placement new subtasks created on this thread will use.
pub fn placement() -> SubtaskPlacement {
    SUBTASK_PLACEMENT_CONTEXT
        .with(|slot| *slot.borrow())
        .unwrap_or_else(placement_from_environment)
}

/// Run `inner` with `placement` installed for every subtask its execution tree
/// creates. The override rides [`AmbientExecutionScope`], so it survives the
/// awaits and thread migrations inside `inner`.
pub fn scope_placement<F: Future>(
    placement: SubtaskPlacement,
    inner: F,
) -> impl Future<Output = F::Output> {
    let mut scope = AmbientExecutionScope::capture_for_inline_subtask();
    scope.set_subtask_placement(Some(placement));
    scope_ambient(scope, inner)
}

/// Wrap a child-interpreter body so it carries its parent's ambient execution
/// scope and pool registry.
///
/// Capture happens here, synchronously, while the creating task's scope is
/// still swapped in. A subtask created from inside a fan-out worker whose own
/// scope has already been swapped out would otherwise read an empty or sibling
/// context: subtasks are independent runtime tasks, not nested inside the
/// parent's poll.
pub(crate) fn prepare<F: Future>(
    registry: Arc<PoolRegistry>,
    future: F,
) -> PreparedSubtask<impl Future<Output = F::Output>> {
    PreparedSubtask {
        inner: scope_ambient(
            AmbientExecutionScope::capture_for_inline_subtask(),
            with_pool_registry_scope(registry, future),
        ),
    }
}

/// Put a prepared subtask on the runtime.
pub(crate) fn spawn<F>(future: PreparedSubtask<F>) -> tokio::task::JoinHandle<F::Output>
where
    F: Future + Send + 'static,
    F::Output: Send + 'static,
{
    match placement() {
        SubtaskPlacement::Worker => tokio::spawn(future),
        SubtaskPlacement::CurrentThread => tokio::task::spawn_local(future),
    }
}

/// Put a prepared subtask on the runtime as a member of `set`.
pub(crate) fn spawn_into<F>(
    set: &mut tokio::task::JoinSet<F::Output>,
    future: PreparedSubtask<F>,
) -> tokio::task::AbortHandle
where
    F: Future + Send + 'static,
    F::Output: Send + 'static,
{
    match placement() {
        SubtaskPlacement::Worker => set.spawn(future),
        SubtaskPlacement::CurrentThread => {
            // A current-thread execution tree is the deterministic placement.
            // Give each child its first poll in source order before admitting
            // the next child; Tokio's local ready queue does not promise the
            // same order across separately constructed runtimes. A pending
            // future is polled again immediately by the spawned task, replacing
            // this noop waker with its real scheduler waker.
            let mut future = Box::pin(future);
            let mut context = Context::from_waker(std::task::Waker::noop());
            match future.as_mut().poll(&mut context) {
                Poll::Ready(value) => set.spawn_local(async move { value }),
                Poll::Pending => set.spawn_local(future),
            }
        }
    }
}

/// Prepare and spawn in one step, for callers that do not hold the future
/// between the two.
pub(crate) fn spawn_child<F>(
    registry: Arc<PoolRegistry>,
    future: F,
) -> tokio::task::JoinHandle<F::Output>
where
    F: Future + Send + 'static,
    F::Output: Send + 'static,
{
    spawn(prepare(registry, future))
}

/// Run lifecycle cleanup independently of the caller's Tokio runtime.
///
/// VM destruction is synchronous and may coincide with embedding-runtime
/// shutdown. Cleanup spawned on that runtime can therefore be cancelled
/// before its first poll. This process-owned runtime is deliberately tiny and
/// exists only for terminal persistence and release of task-owned sessions.
pub(crate) fn spawn_lifecycle_cleanup<F>(future: F) -> tokio::task::JoinHandle<F::Output>
where
    F: Future + Send + 'static,
    F::Output: Send + 'static,
{
    static CLEANUP_RUNTIME: std::sync::OnceLock<tokio::runtime::Runtime> =
        std::sync::OnceLock::new();
    let runtime = CLEANUP_RUNTIME.get_or_init(|| {
        tokio::runtime::Builder::new_multi_thread()
            .thread_stack_size(crate::RUNTIME_STACK_SIZE)
            .worker_threads(1)
            .thread_name("harn-lifecycle-cleanup")
            .enable_all()
            .build()
            .expect("build Harn lifecycle cleanup runtime")
    });
    note_lifecycle_cleanup_spawn();
    runtime.spawn(future)
}

#[cfg(test)]
thread_local! {
    /// Detached recovery tasks this thread has handed to the process-owned
    /// runtime.
    ///
    /// Whether an execution transfers cleanup is decided synchronously as it
    /// drops, so a count is the exact observable for "this drop scheduled no
    /// recovery" — a claim no amount of waiting can establish. It is counted
    /// per thread rather than per process because the decision happens on the
    /// dropping thread: a process-wide total answers "did anything anywhere
    /// schedule recovery", which a sibling case sharing the test binary can
    /// move under the reader's feet, and a `before == after` assertion then
    /// fails having observed someone else's work (harn#7960).
    static LIFECYCLE_CLEANUP_SPAWNS: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
}

#[cfg(test)]
fn note_lifecycle_cleanup_spawn() {
    LIFECYCLE_CLEANUP_SPAWNS.with(|count| count.set(count.get().saturating_add(1)));
}

#[cfg(not(test))]
fn note_lifecycle_cleanup_spawn() {}

#[cfg(test)]
pub(crate) fn lifecycle_cleanup_spawn_count() -> u64 {
    LIFECYCLE_CLEANUP_SPAWNS.with(std::cell::Cell::get)
}

/// Spawn a long-lived child that inherits execution policy but owns an
/// independent session lifecycle.
pub(crate) fn spawn_inherited_child<F>(
    registry: Arc<PoolRegistry>,
    future: F,
) -> tokio::task::JoinHandle<F::Output>
where
    F: Future + Send + 'static,
    F::Output: Send + 'static,
{
    spawn(PreparedSubtask {
        inner: scope_ambient(
            AmbientExecutionScope::capture_inherited(),
            with_pool_registry_scope(registry, future),
        ),
    })
}

#[cfg(test)]
#[path = "subtask/cross_thread_tests.rs"]
mod cross_thread_tests;

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

    #[test]
    fn placement_names_round_trip() {
        assert_eq!(
            SubtaskPlacement::from_env_value("worker"),
            Ok(SubtaskPlacement::Worker)
        );
        assert_eq!(
            SubtaskPlacement::from_env_value(" CURRENT_THREAD "),
            Ok(SubtaskPlacement::CurrentThread)
        );
        assert_eq!(
            SubtaskPlacement::from_env_value("sideways")
                .expect_err("invalid placement must not become an absent override")
                .to_string(),
            "invalid HARN_VM_SUBTASK_PLACEMENT value \"sideways\"; expected one of worker, current_thread"
        );
        assert_eq!(SubtaskPlacement::Worker.to_string(), "worker");
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn scoped_placement_reaches_the_spawn_seam() {
        assert_eq!(placement(), SubtaskPlacement::Worker);
        let observed =
            scope_placement(SubtaskPlacement::CurrentThread, async { placement() }).await;
        assert_eq!(observed, SubtaskPlacement::CurrentThread);
        assert_eq!(placement(), SubtaskPlacement::Worker);
    }

    #[test]
    fn placement_selects_the_executor_thread() {
        fn observed_thread(
            runtime: &tokio::runtime::Runtime,
            placement: SubtaskPlacement,
        ) -> std::thread::ThreadId {
            runtime.block_on(async {
                tokio::task::LocalSet::new()
                    .run_until(scope_placement(placement, async {
                        spawn(PreparedSubtask {
                            inner: async { std::thread::current().id() },
                        })
                        .await
                        .expect("subtask completes")
                    }))
                    .await
            })
        }

        let runtime = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(2)
            .enable_all()
            .build()
            .expect("test runtime");
        let creating_thread = std::thread::current().id();

        assert_ne!(
            observed_thread(&runtime, SubtaskPlacement::Worker),
            creating_thread
        );
        assert_eq!(
            observed_thread(&runtime, SubtaskPlacement::CurrentThread),
            creating_thread
        );
    }
}