arty_executor 0.5.0

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

use observed::Sink;
use observed::context::Transfer;

use crate::telemetry;

/// Describes how the submitter selected this executor for a task.
///
/// This value is lifecycle telemetry metadata. It does not change how the executor schedules or
/// polls the task.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum TaskPlacement {
    /// The submitter did not require a specific worker.
    #[default]
    Any,

    /// The submitter targeted its current worker.
    SameThread,
}

/// Lifecycle metadata for a task registered with [`TaskSet`](crate::TaskSet).
///
/// The executor applies captured observed context while polling, cleaning up, and emitting
/// lifecycle events so all task activity keeps the submission dimensions. Placement records how
/// the caller chose the executor; it does not select a worker. The default options add no context
/// and report [`TaskPlacement::Any`].
#[derive(Debug, Default)]
pub struct TaskOptions {
    context: Option<Transfer>,
    placement: TaskPlacement,
}

pub(crate) struct TaskCreatedReporter {
    context: Option<Transfer>,
    placement: TaskPlacement,
    sink: Sink,
}

impl TaskCreatedReporter {
    pub(crate) fn report(self) {
        if let Some(context) = self.context {
            let _guard = context.apply_current_thread();
            telemetry::report_created(&self.sink, self.placement);
        } else {
            telemetry::report_created(&self.sink, self.placement);
        }
    }
}

impl TaskOptions {
    /// Creates options with the specified placement metadata.
    #[must_use]
    pub const fn new(placement: TaskPlacement) -> Self {
        Self { context: None, placement }
    }

    /// Creates options that apply the captured observed `context` throughout the task lifecycle.
    #[must_use]
    pub fn observed(context: Transfer, placement: TaskPlacement) -> Self {
        Self {
            context: Some(context),
            placement,
        }
    }

    pub(crate) fn with_context<T>(&self, f: impl FnOnce() -> T) -> T {
        let Some(context) = &self.context else {
            return f();
        };

        let _guard = context.apply_current_thread();
        f()
    }

    pub(crate) fn created_reporter(&self, sink: Sink) -> TaskCreatedReporter {
        TaskCreatedReporter {
            context: self.context.clone(),
            placement: self.placement,
            sink,
        }
    }

    pub(crate) fn report_succeeded(&self, sink: &Sink, polls: u64) {
        self.with_context(|| telemetry::report_succeeded(sink, polls));
    }

    pub(crate) fn report_panicked(&self, sink: &Sink, polls: u64) {
        self.with_context(|| telemetry::report_panicked(sink, polls));
    }
}

#[cfg(test)]
#[cfg_attr(coverage_nightly, coverage(off))]
mod tests {
    use std::panic::{RefUnwindSafe, UnwindSafe};

    use static_assertions::{assert_impl_all, assert_not_impl_any};

    use super::*;

    assert_impl_all!(TaskPlacement: UnwindSafe, RefUnwindSafe);
    // Captured enrichment can contain user-defined interior state, so options intentionally do
    // not promise that they may cross a caught-unwind boundary.
    assert_not_impl_any!(TaskOptions: UnwindSafe, RefUnwindSafe);

    #[test]
    fn observed_options_keep_context_and_placement() {
        let options = TaskOptions::observed(Sink::noop().transfer_context(), TaskPlacement::SameThread);

        assert!(options.context.is_some());
        assert_eq!(options.placement, TaskPlacement::SameThread);
    }

    #[test]
    fn new_options_keep_placement_without_context() {
        let options = TaskOptions::new(TaskPlacement::SameThread);

        assert!(options.context.is_none());
        assert_eq!(options.placement, TaskPlacement::SameThread);
    }
}