otel-arrow-dfe-engine 0.61.0

// Copyright The OpenTelemetry Authors
// SPDX-License-Identifier: Apache-2.0

//! Thread/task-local storage for entity keys used by our internal telemetry instrumentation
//! to associate metrics and events with the correct pipeline entity, node entity or the correct
//! input/output channel entities.

use otel_arrow_dfe_config::PortName;
use otel_arrow_dfe_telemetry::metrics::{
    MeasurementMetricSet, MeasurementMetricSetHandler, MetricSet, MetricSetHandler,
};
use otel_arrow_dfe_telemetry::registry::{EntityKey, MetricSetKey, TelemetryRegistryHandle};
use std::cell::{Cell, RefCell};
use std::fmt::Debug;
use std::future::Future;
use std::rc::Rc;

thread_local! {
    /// Defined when building a pipeline to associate telemetry with the pipeline entity.
    /// Present for the pipeline lifetime.
    static PIPELINE_ENTITY_KEY: Cell<Option<EntityKey>> = const { Cell::new(None) };
    /// Set for each node being built to provide telemetry handle during construction.
    static BUILD_NODE_TELEMETRY_HANDLE: RefCell<Option<NodeTelemetryHandle>> = const { RefCell::new(None) };
}

/// RAII guard that clears and unregisters the pipeline entity on drop.
pub struct PipelineEntityScope {
    registry: TelemetryRegistryHandle,
    key: EntityKey,
}

impl Drop for PipelineEntityScope {
    fn drop(&mut self) {
        PIPELINE_ENTITY_KEY.with(|cell| cell.set(None));
        let _ = self.registry.unregister_entity(self.key);
    }
}

/// Sets the pipeline entity key for the current thread.
#[must_use]
pub fn set_pipeline_entity_key(
    registry: TelemetryRegistryHandle,
    key: EntityKey,
) -> PipelineEntityScope {
    PIPELINE_ENTITY_KEY.with(|cell| cell.set(Some(key)));
    PipelineEntityScope { registry, key }
}

/// Returns the pipeline entity key for the current thread, if set.
#[must_use]
pub fn pipeline_entity_key() -> Option<EntityKey> {
    PIPELINE_ENTITY_KEY.with(Cell::get)
}

tokio::task_local! {
    /// Task-local context for nodes, including entity keys and telemetry handle.
    /// This task local can be used by our internal telemetry instrumentation to associate metrics
    /// and events with the correct node entity or the correct input/output channels.
    static NODE_TASK_CONTEXT: NodeTaskContext;
}

/// Per-task snapshot for fast lookups.
/// Separate from NodeTelemetryState to avoid RefCell borrows on hot paths during node execution.
#[derive(Debug)]
pub(crate) struct NodeTaskContext {
    telemetry_handle: Option<NodeTelemetryHandle>,

    /// Entities associated with this node task.
    entity_key: Option<EntityKey>,
    input_channel_key: Option<EntityKey>,
    output_channel_keys: Vec<(PortName, EntityKey)>,
}

impl NodeTaskContext {
    pub(crate) fn new(
        entity_key: Option<EntityKey>,
        telemetry_handle: Option<NodeTelemetryHandle>,
        input_channel_key: Option<EntityKey>,
        output_channel_keys: Vec<(PortName, EntityKey)>,
    ) -> Self {
        Self {
            entity_key,
            telemetry_handle,
            input_channel_key,
            output_channel_keys,
        }
    }
}

/// Returns the node entity key for the current task, if set.
#[inline]
#[must_use]
pub fn node_entity_key() -> Option<EntityKey> {
    NODE_TASK_CONTEXT
        .try_with(|ctx| ctx.entity_key)
        .ok()
        .flatten()
}

/// Returns the input channel entity key for the current task, if set.
#[inline]
#[must_use]
pub fn node_input_channel_key() -> Option<EntityKey> {
    NODE_TASK_CONTEXT
        .try_with(|ctx| ctx.input_channel_key)
        .ok()
        .flatten()
}

/// Returns the output channel entity key for the given port in the current task, if set.
///
/// Implementation detail: This function looks up the output channel keys stored in the
/// NodeTaskContext for the current task, searching for a matching port name and returning
/// the associated EntityKey if found. If the number of output channels becomes large, consider
/// optimizing the storage structure for faster lookups. Right now, a linear search seems a good
/// trade-off given the expected small number of output channels per node.
#[inline]
#[must_use]
pub fn node_output_channel_key(port: &str) -> Option<EntityKey> {
    NODE_TASK_CONTEXT
        .try_with(|ctx| {
            ctx.output_channel_keys
                .iter()
                .find(|(name, _)| name.as_ref() == port)
                .map(|(_, key)| *key)
        })
        .ok()
        .flatten()
}

/// Runs the given future with the provided node task context in task-local storage.
pub(crate) fn instrument_with_node_context<F, T>(
    ctx: NodeTaskContext,
    fut: F,
) -> impl Future<Output = T>
where
    F: Future<Output = T>,
{
    NODE_TASK_CONTEXT.scope(ctx, fut)
}

/// Returns the current node telemetry handle, if set.
pub(crate) fn current_node_telemetry_handle() -> Option<NodeTelemetryHandle> {
    // Runtime code uses the task-local context.
    // Build-time code uses the thread-local fallback before node tasks are spawned.
    if let Ok(handle) = NODE_TASK_CONTEXT.try_with(|ctx| ctx.telemetry_handle.clone()) {
        return handle;
    }
    BUILD_NODE_TELEMETRY_HANDLE.with(|cell| cell.borrow().clone())
}

/// Runs a closure with the provided node telemetry handle in thread-local storage.
pub(crate) fn with_node_telemetry_handle<T>(
    handle: NodeTelemetryHandle,
    f: impl FnOnce() -> T,
) -> T {
    // Build-time helper: node construction happens before task locals exist, so we scope a
    // thread-local handle to let builders access current_node_telemetry_handle().
    BUILD_NODE_TELEMETRY_HANDLE.with(|cell| {
        let _ = cell.replace(Some(handle));
        let result = f();
        let _ = cell.replace(None);
        result
    })
}

/// Scope-agnostic per-entity telemetry bookkeeping: the registered
/// `EntityKey`, metric sets tied to it, and extra entities whose lifetime is
/// bound to it. Usable for any registered entity (pipeline, node, extension,
/// channel, ...) -- the kind is determined by which attribute set was used to
/// mint the `EntityKey`.
#[derive(Clone)]
pub(crate) struct EntityTelemetryHandle {
    registry: TelemetryRegistryHandle,
    state: Rc<RefCell<EntityTelemetryState>>,
}

impl Debug for EntityTelemetryHandle {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("EntityTelemetryHandle")
            .field("entity_key", &self.state.borrow().entity_key)
            .finish()
    }
}

struct EntityTelemetryState {
    entity_key: EntityKey,
    metric_keys: Vec<MetricSetKey>,
    extra_entity_keys: Vec<EntityKey>,
    cleaned: bool,
}

impl EntityTelemetryHandle {
    pub(crate) fn new(registry: TelemetryRegistryHandle, entity_key: EntityKey) -> Self {
        Self {
            registry,
            state: Rc::new(RefCell::new(EntityTelemetryState {
                entity_key,
                metric_keys: Vec::new(),
                extra_entity_keys: Vec::new(),
                cleaned: false,
            })),
        }
    }

    pub(crate) fn entity_key(&self) -> EntityKey {
        self.state.borrow().entity_key
    }

    pub(crate) fn register_metric_set<T: MetricSetHandler + Default + Debug + Send + Sync>(
        &self,
    ) -> MetricSet<T> {
        let entity_key = self.state.borrow().entity_key;
        self.register_metric_set_for_entity::<T>(entity_key)
    }

    /// Registers a metric set for an arbitrary entity key and tracks the resulting
    /// metric-set key on this handle so it is unregistered as part of cleanup.
    ///
    /// Used when an extension or other host attaches metrics to a child entity
    /// (e.g. a control channel) whose lifecycle is tied to this handle.
    pub(crate) fn register_metric_set_for_entity<
        T: MetricSetHandler + Default + Debug + Send + Sync,
    >(
        &self,
        entity_key: EntityKey,
    ) -> MetricSet<T> {
        let metrics = self
            .registry
            .register_metric_set_for_entity::<T>(entity_key);
        self.track_metric_set(metrics.metric_set_key());
        metrics
    }

    /// Registers a measurement metric set for an arbitrary entity and tracks it for cleanup.
    pub(crate) fn register_measurement_metric_set_for_entity<
        T: MeasurementMetricSetHandler + Debug + Send + Sync,
    >(
        &self,
        entity_key: EntityKey,
    ) -> MeasurementMetricSet<T> {
        let metrics = self
            .registry
            .register_metric_set_with_measurement_attributes_for_entity::<T>(entity_key);
        self.track_metric_set(metrics.metric_set_key());
        metrics
    }

    pub(crate) fn track_metric_set(&self, metrics_key: MetricSetKey) {
        self.state.borrow_mut().metric_keys.push(metrics_key);
    }

    pub(crate) fn track_entity(&self, entity_key: EntityKey) {
        self.state.borrow_mut().extra_entity_keys.push(entity_key);
    }

    /// Unregister tracked metric sets, extra entities, and the entity itself.
    /// Idempotent.
    pub(crate) fn cleanup(&self) {
        let mut state = self.state.borrow_mut();
        if state.cleaned {
            return;
        }
        state.cleaned = true;
        let metric_keys = std::mem::take(&mut state.metric_keys);
        let extra_entity_keys = std::mem::take(&mut state.extra_entity_keys);
        let entity_key = state.entity_key;
        drop(state);

        for key in metric_keys {
            let _ = self.registry.unregister_metric_set(key);
        }
        for key in extra_entity_keys {
            let _ = self.registry.unregister_entity(key);
        }
        let _ = self.registry.unregister_entity(entity_key);
    }
}

/// RAII guard that cleans up an `EntityTelemetryHandle` on drop. Scope-agnostic.
#[doc(hidden)]
pub struct EntityTelemetryGuard {
    handle: EntityTelemetryHandle,
}

impl EntityTelemetryGuard {
    pub(crate) const fn new(handle: EntityTelemetryHandle) -> Self {
        Self { handle }
    }
}

impl Drop for EntityTelemetryGuard {
    fn drop(&mut self) {
        self.handle.cleanup();
    }
}

// Per-node channel-port wiring. Node-only add-on layered on top of EntityTelemetryHandle.
#[derive(Default)]
struct NodeChannelState {
    input_channel_key: Option<EntityKey>,
    output_channel_keys: Vec<(PortName, EntityKey)>,
    control_channel_key: Option<EntityKey>,
}

/// Handle for per-node telemetry state: entity + metrics (via the embedded
/// `EntityTelemetryHandle`) plus node-only channel-port wiring.
#[derive(Clone)]
pub(crate) struct NodeTelemetryHandle {
    entity: EntityTelemetryHandle,
    channels: Rc<RefCell<NodeChannelState>>,
}

impl Debug for NodeTelemetryHandle {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("NodeTelemetryHandle")
            .field("entity_key", &self.entity.entity_key())
            .finish()
    }
}

impl NodeTelemetryHandle {
    pub(crate) fn new(registry: TelemetryRegistryHandle, entity_key: EntityKey) -> Self {
        Self {
            entity: EntityTelemetryHandle::new(registry, entity_key),
            channels: Rc::new(RefCell::new(NodeChannelState::default())),
        }
    }

    pub(crate) fn entity_key(&self) -> EntityKey {
        self.entity.entity_key()
    }

    #[allow(dead_code)]
    pub(crate) fn entity_handle(&self) -> EntityTelemetryHandle {
        self.entity.clone()
    }

    pub(crate) fn register_metric_set<T: MetricSetHandler + Default + Debug + Send + Sync>(
        &self,
    ) -> MetricSet<T> {
        self.entity.register_metric_set::<T>()
    }

    pub(crate) fn track_metric_set(&self, metrics_key: MetricSetKey) {
        self.entity.track_metric_set(metrics_key);
    }

    pub(crate) fn track_entity(&self, entity_key: EntityKey) {
        self.entity.track_entity(entity_key);
    }

    pub(crate) fn set_input_channel_key(&self, key: EntityKey) {
        let mut state = self.channels.borrow_mut();
        debug_assert!(
            state.input_channel_key.is_none(),
            "input channel key already set"
        );
        state.input_channel_key = Some(key);
    }

    pub(crate) fn add_output_channel_key(&self, port: PortName, key: EntityKey) {
        let mut state = self.channels.borrow_mut();
        if let Some((_, existing_key)) = state
            .output_channel_keys
            .iter_mut()
            .find(|(name, _)| name == &port)
        {
            *existing_key = key;
            return;
        }
        state.output_channel_keys.push((port, key));
    }

    pub(crate) fn set_control_channel_key(&self, key: EntityKey) {
        let mut state = self.channels.borrow_mut();
        debug_assert!(
            state.control_channel_key.is_none(),
            "control channel key already set"
        );
        state.control_channel_key = Some(key);
    }

    pub(crate) fn input_channel_key(&self) -> Option<EntityKey> {
        self.channels.borrow().input_channel_key
    }

    pub(crate) fn output_channel_keys(&self) -> Vec<(PortName, EntityKey)> {
        self.channels.borrow().output_channel_keys.clone()
    }

    /// Unregister channel entities, then delegate to the entity handle for
    /// metric-set + entity cleanup. Idempotent.
    pub(crate) fn cleanup(&self) {
        let (input, control, outputs) = {
            let mut state = self.channels.borrow_mut();
            (
                state.input_channel_key.take(),
                state.control_channel_key.take(),
                std::mem::take(&mut state.output_channel_keys),
            )
        };
        if let Some(key) = input {
            let _ = self.entity.registry.unregister_entity(key);
        }
        if let Some(key) = control {
            let _ = self.entity.registry.unregister_entity(key);
        }
        for (_, key) in outputs {
            let _ = self.entity.registry.unregister_entity(key);
        }
        self.entity.cleanup();
    }
}

#[doc(hidden)]
pub struct NodeTelemetryGuard {
    handle: NodeTelemetryHandle,
}

impl NodeTelemetryGuard {
    pub(crate) const fn new(handle: NodeTelemetryHandle) -> Self {
        Self { handle }
    }

    pub(crate) fn entity_key(&self) -> EntityKey {
        self.handle.entity_key()
    }

    pub(crate) fn handle(&self) -> NodeTelemetryHandle {
        self.handle.clone()
    }
}

impl Drop for NodeTelemetryGuard {
    fn drop(&mut self) {
        self.handle.cleanup();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::attributes::{ChannelImplementation, ChannelKind, ChannelMode, ChannelType};
    use crate::channel_metrics::{ChannelReceiverMetrics, ChannelSenderMetrics};
    use crate::context::ControllerContext;
    use crate::pipeline_metrics::PipelineMetricsMonitor;
    use otel_arrow_dfe_config::node::NodeKind;
    use otel_arrow_dfe_telemetry::registry::TelemetryRegistryHandle;
    use std::borrow::Cow;
    use std::collections::HashMap;

    /// Scenario: A pipeline with node and channel telemetry finishes and drops every guard.
    /// Guarantees: All registered entities and metric sets, including measurement sets, are removed.
    #[test]
    fn pipeline_cleanup_unregisters_entities() {
        let registry = TelemetryRegistryHandle::new();
        let controller_ctx = ControllerContext::new(registry.clone());
        let pipeline_ctx =
            controller_ctx.pipeline_context_with("group".into(), "pipe".into(), 0, 1, 0);

        let pipeline_entity_key = pipeline_ctx.register_pipeline_entity();
        let _pipeline_entity_guard =
            set_pipeline_entity_key(pipeline_ctx.metrics_registry(), pipeline_entity_key);
        let _pipeline_metrics = PipelineMetricsMonitor::new(pipeline_ctx.clone());

        let source_ctx = pipeline_ctx.with_node_context(
            "source".into(),
            "urn:test:receiver:example".into(),
            NodeKind::Receiver,
            HashMap::new(),
        );
        let dest_ctx = pipeline_ctx.with_node_context(
            "dest".into(),
            "urn:test:processor:example".into(),
            NodeKind::Processor,
            HashMap::new(),
        );

        let source_entity_key = source_ctx.register_node_entity();
        let dest_entity_key = dest_ctx.register_node_entity();
        let source_handle =
            NodeTelemetryHandle::new(source_ctx.metrics_registry(), source_entity_key);
        let dest_handle = NodeTelemetryHandle::new(dest_ctx.metrics_registry(), dest_entity_key);
        let source_guard = NodeTelemetryGuard::new(source_handle.clone());
        let dest_guard = NodeTelemetryGuard::new(dest_handle.clone());

        let channel_id: Cow<'static, str> = "chan:pdata".into();
        let out_key = source_ctx.register_node_channel_entity(
            channel_id.clone(),
            "out".into(),
            ChannelKind::Pdata,
            ChannelMode::Local,
            ChannelType::Mpsc,
            ChannelImplementation::Internal,
        );
        source_handle.add_output_channel_key("out".into(), out_key);
        let in_key = dest_ctx.register_node_channel_entity(
            channel_id,
            "input".into(),
            ChannelKind::Pdata,
            ChannelMode::Local,
            ChannelType::Mpsc,
            ChannelImplementation::Internal,
        );
        dest_handle.set_input_channel_key(in_key);
        let ctrl_key = source_ctx.register_node_channel_entity(
            "chan:ctrl".into(),
            "input".into(),
            ChannelKind::Control,
            ChannelMode::Local,
            ChannelType::Mpsc,
            ChannelImplementation::Internal,
        );
        source_handle.set_control_channel_key(ctrl_key);

        let out_metrics =
            source_ctx.register_measurement_metric_set_for_entity::<ChannelSenderMetrics>(out_key);
        source_handle.track_metric_set(out_metrics.metric_set_key());
        let in_metrics =
            dest_ctx.register_measurement_metric_set_for_entity::<ChannelReceiverMetrics>(in_key);
        dest_handle.track_metric_set(in_metrics.metric_set_key());
        let _ = source_handle
            .entity_handle()
            .register_measurement_metric_set_for_entity::<ChannelSenderMetrics>(source_entity_key);
        let _ = dest_handle
            .entity_handle()
            .register_measurement_metric_set_for_entity::<ChannelSenderMetrics>(dest_entity_key);

        assert_eq!(registry.entity_count(), 6);
        assert_eq!(registry.metric_set_count(), 6);

        drop(dest_guard);
        drop(source_guard);
        drop(_pipeline_metrics);
        drop(_pipeline_entity_guard);

        assert_eq!(registry.metric_set_count(), 0);
        assert_eq!(registry.entity_count(), 0);
    }

    /// Scenario: Multiple owners request cleanup through the same node telemetry handle.
    /// Guarantees: Repeated cleanup is a no-op and never unregisters a reused entity key.
    #[test]
    fn entity_telemetry_handle_cleanup_is_idempotent() {
        // Handle state is `Rc<RefCell<...>>` so multiple paths can call
        // `cleanup`. Second call must be a no-op; otherwise we'd
        // double-unregister an entity whose key may have been reissued.
        let registry = TelemetryRegistryHandle::new();
        let controller_ctx = ControllerContext::new(registry.clone());
        let pipeline_ctx =
            controller_ctx.pipeline_context_with("group".into(), "pipe".into(), 0, 1, 0);
        let ext_ctx = pipeline_ctx.extension_context();

        let baseline_entities = registry.entity_count();
        let baseline_metric_sets = registry.metric_set_count();

        let entity_key = ext_ctx.register_extension_entity(
            "idem".into(),
            crate::extension::wrapper::ExtensionVariant::Local,
        );
        let handle = EntityTelemetryHandle::new(ext_ctx.metrics_registry(), entity_key);

        let extra_entity = ext_ctx.register_extension_channel_entity(
            "idem".into(),
            crate::extension::wrapper::ExtensionVariant::Local,
            "ctrl".into(),
            ChannelMode::Local,
            ChannelImplementation::Internal,
        );
        handle.track_entity(extra_entity);
        let _metrics =
            handle.register_measurement_metric_set_for_entity::<ChannelSenderMetrics>(entity_key);

        assert_eq!(registry.entity_count(), baseline_entities + 2);
        assert_eq!(registry.metric_set_count(), baseline_metric_sets + 1);

        let clone = handle.clone();
        clone.cleanup();
        assert_eq!(registry.entity_count(), baseline_entities);
        assert_eq!(registry.metric_set_count(), baseline_metric_sets);

        // Re-register to expose any latent double-unregister.
        let resurrected = ext_ctx.register_extension_entity(
            "resurrected".into(),
            crate::extension::wrapper::ExtensionVariant::Local,
        );
        assert_eq!(registry.entity_count(), baseline_entities + 1);

        let guard = EntityTelemetryGuard::new(handle);
        drop(guard);
        assert_eq!(
            registry.entity_count(),
            baseline_entities + 1,
            "second cleanup must not unregister the resurrected entity"
        );

        let _ = registry.unregister_entity(resurrected);
    }
}