otel-arrow-dfe-engine 0.61.0

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

//! Context providing general information on the current controller and the current pipeline.

use crate::Interests;
use crate::attributes::{
    ChannelImplementation, ChannelKind, ChannelMode, ChannelType, CustomAttributeSet,
    EngineAttributeSet, EngineEntityAttributeSet, ExtensionAttributeSet,
    ExtensionChannelAttributeSet, ExtensionScopeAttributeSet, NodeChannelAttributeSet,
    NodeWithCustomChannelAttributeSet, NodeWithCustomTopicAttributeSet, NodeWithTopicAttributeSet,
    PipelineAttributeSet, config_map_to_telemetry,
};
pub use crate::attributes::{NodeAttributeSet, NodeWithCustomAttributeSet};
use crate::context_declaration::CompiledContextBindings;
use crate::entity_context::{current_node_telemetry_handle, node_entity_key};
use crate::listener_group::ListenerGroupSnapshot;
use crate::memory_limiter::MemoryPressureState;
use crate::node::NodeId as EngineNodeId;
use data_encoding::BASE32_NOPAD;
use otel_arrow_dfe_config::node::NodeKind;
use otel_arrow_dfe_config::pipeline::telemetry::TelemetryAttribute;
use otel_arrow_dfe_config::policy::DistributionTier;
use otel_arrow_dfe_config::{
    NodeId as ConfigNodeId, NodeUrn, PipelineGroupId, PipelineId, PipelineKey,
};
use otel_arrow_dfe_telemetry::InternalTelemetrySettings;
use otel_arrow_dfe_telemetry::attributes::AttributeSetHandler;
use otel_arrow_dfe_telemetry::metrics::MetricSetRegistrar;
use otel_arrow_dfe_telemetry::metrics::{
    MeasurementMetricSet, MeasurementMetricSetHandler, MetricSet, MetricSetHandler,
    RegistrationMetricSetHandler,
};
use otel_arrow_dfe_telemetry::registry::{EntityKey, MetricSetKey, TelemetryRegistryHandle};
use std::any::Any;
use std::borrow::Cow;
use std::collections::HashMap;
use std::fmt::Debug;
use std::sync::{Arc, LazyLock};
use uuid::Uuid;

/// A shared, immutable mapping from otel_arrow_dfe_config node names
/// (without index numbers) to their engine-specific pipeline indices.
pub type NodeNameIndex = Arc<HashMap<ConfigNodeId, EngineNodeId>>;

// Generate a stable, unique identifier per process instance (base32-encoded UUID v7).
static PROCESS_INSTANCE_ID: LazyLock<Cow<'static, str>> = LazyLock::new(|| {
    let uuid = Uuid::now_v7();
    Cow::Owned(BASE32_NOPAD.encode(uuid.as_bytes()))
});

fn detect_host_id() -> Option<String> {
    if let Ok(host) = std::env::var("HOSTNAME")
        && !host.is_empty()
    {
        return Some(host);
    }
    std::fs::read_to_string("/etc/hostname")
        .ok()
        .map(|host| host.trim().to_owned())
        .filter(|host| !host.is_empty())
}

fn detect_container_id() -> Option<String> {
    let cgroup = std::fs::read_to_string("/proc/self/cgroup").ok()?;
    for line in cgroup.lines() {
        let path = line.split(':').nth(2).unwrap_or("");
        for part in path.split('/') {
            let token = part.trim();
            if (32..=128).contains(&token.len())
                && token
                    .chars()
                    .all(|c| c.is_ascii_hexdigit() || matches!(c, '.' | '-' | '_'))
                && token.chars().filter(|c| c.is_ascii_hexdigit()).count() >= 32
            {
                return Some(token.to_owned());
            }
        }
    }
    None
}

static HOST_ID: LazyLock<Cow<'static, str>> =
    LazyLock::new(|| detect_host_id().map_or(Cow::Borrowed(""), Cow::Owned));
static CONTAINER_ID: LazyLock<Cow<'static, str>> =
    LazyLock::new(|| detect_container_id().map_or(Cow::Borrowed(""), Cow::Owned));

/// A lightweight/cloneable controller context.
#[derive(Clone, Debug)]
pub struct ControllerContext {
    telemetry_registry_handle: TelemetryRegistryHandle,
    /// Unique process instance identifier (base32-encoded UUID v7).
    process_instance_id: Cow<'static, str>,
    /// Host identifier, when available (e.g. hostname).
    host_id: Cow<'static, str>,
    /// Container identifier, when available (e.g. Docker or containerd container ID).
    container_id: Cow<'static, str>,
    memory_pressure_state: MemoryPressureState,
    state_directory: Option<crate::state_dir::StateDirectory>,
}

/// Parameters required to create a pipeline context.
#[derive(Clone, Debug)]
pub struct PipelineContextParams {
    /// Pipeline group ID for the current pipeline execution context.
    pub pipeline_group_id: PipelineGroupId,
    /// Pipeline ID for the current pipeline execution context.
    pub pipeline_id: PipelineId,
    /// Core ID for the current pipeline execution context.
    pub core_id: usize,
    /// Total number of cores allocated to this pipeline.
    /// Used by nodes that need to share resources across cores (e.g., disk budgets).
    pub num_cores: usize,
    /// Thread ID for the current pipeline execution context.
    pub thread_id: usize,
    /// NUMA node ID resolved for the current pipeline execution context.
    pub numa_node_id: usize,
}

/// A lightweight/cloneable pipeline context.
#[derive(Clone, Debug)]
pub struct PipelineContext {
    controller_context: ControllerContext,
    pipeline_context_params: PipelineContextParams,
    deployment_generation: u64,
    pipeline_telemetry_attrs: HashMap<String, TelemetryAttribute>,
    node_id: ConfigNodeId,
    node_urn: NodeUrn,
    node_kind: NodeKind,
    node_interests: Interests,
    node_duration_distribution: DistributionTier,
    node_telemetry_attrs: HashMap<String, TelemetryAttribute>,
    admission: crate::admission::AdmissionBinder,

    /// Internal telemetry settings for the Internal Telemetry Receiver (ITR).
    /// Only the ITR factory reads this; other receivers ignore it.
    internal_telemetry: Option<InternalTelemetrySettings>,
    /// Shared mapping from node names to pipeline indices for mapping
    /// node names to the index used to send node control messages by,
    /// for example to map source-node name to index for inferring
    /// routes at runtime (e.g., how crates/validation works).
    node_names: NodeNameIndex,
    /// Optional pipeline-scoped topic set injected by the controller.
    /// ToDo: Make PipelineContext generic over a TopicSet type to avoid dynamic typing here.
    topic_set: Option<Arc<dyn Any + Send + Sync>>,
    // Immutable placement metadata shared across contexts for this pipeline instance.
    // Consumers should cache any needed listener plan during setup rather than cloning
    // or searching this snapshot from the per-record data path.
    listener_group_snapshot: Arc<ListenerGroupSnapshot>,
    /// Compiled context bindings shared by this runtime's nodes.
    compiled_context_bindings: Arc<CompiledContextBindings>,
}

/// Registrar that binds generated metric-set registration to an existing entity.
///
/// Channel-owned metric sets use this when their scope is an input or output
/// channel rather than the node entity selected by [`PipelineContext`] itself.
pub struct EntityMetricSetRegistrar<'a> {
    pipeline_context: &'a PipelineContext,
    entity_key: EntityKey,
}

impl ControllerContext {
    /// Installs the startup-provisioned root before contexts are distributed.
    #[must_use]
    pub fn with_state_directory(mut self, root: crate::state_dir::StateDirectory) -> Self {
        self.state_directory = Some(root);
        self
    }

    /// Returns the engine [StateDirectory](crate::state_dir::StateDirectory),
    /// or `None` when no state root is configured.
    #[must_use]
    pub fn state_directory(&self) -> Option<&crate::state_dir::StateDirectory> {
        self.state_directory.as_ref()
    }

    /// Creates a new `ControllerContext`.
    #[must_use]
    pub fn new(telemetry_registry_handle: TelemetryRegistryHandle) -> Self {
        Self {
            telemetry_registry_handle,
            process_instance_id: PROCESS_INSTANCE_ID.clone(),
            host_id: HOST_ID.clone(),
            container_id: CONTAINER_ID.clone(),
            memory_pressure_state: MemoryPressureState::default(),
            state_directory: None,
        }
    }

    /// Creates a `ControllerContext` with explicit auto-detected identity values.
    ///
    /// Test-only: the production [`new`](Self::new) constructor derives these
    /// from the host environment, which is unsuitable for asserting the
    /// semantic-convention mapping in [`resource_attributes`](Self::resource_attributes).
    #[cfg(test)]
    fn new_with_identity(
        telemetry_registry_handle: TelemetryRegistryHandle,
        process_instance_id: impl Into<Cow<'static, str>>,
        host_id: impl Into<Cow<'static, str>>,
        container_id: impl Into<Cow<'static, str>>,
    ) -> Self {
        Self {
            telemetry_registry_handle,
            process_instance_id: process_instance_id.into(),
            host_id: host_id.into(),
            container_id: container_id.into(),
            memory_pressure_state: MemoryPressureState::default(),
            state_directory: None,
        }
    }

    /// Returns a new pipeline context with the given identifiers and the current controller context
    /// as the parent context.
    #[must_use]
    pub fn pipeline_context_with(
        &self,
        pipeline_group_id: PipelineGroupId,
        pipeline_id: PipelineId,
        core_id: usize,
        num_cores: usize,
        thread_id: usize,
    ) -> PipelineContext {
        self.pipeline_context_with_generation(
            pipeline_group_id,
            pipeline_id,
            core_id,
            num_cores,
            thread_id,
            0,
        )
    }

    /// Returns a new pipeline context with an explicit deployment generation.
    #[must_use]
    pub fn pipeline_context_with_generation(
        &self,
        pipeline_group_id: PipelineGroupId,
        pipeline_id: PipelineId,
        core_id: usize,
        num_cores: usize,
        thread_id: usize,
        deployment_generation: u64,
    ) -> PipelineContext {
        PipelineContext::new_with_generation(
            self.clone(),
            PipelineContextParams {
                pipeline_group_id,
                pipeline_id,
                core_id,
                num_cores,
                thread_id,
                numa_node_id: 0,
            },
            deployment_generation,
        )
    }

    /// Returns a new pipeline context with explicit NUMA placement metadata.
    #[must_use]
    pub fn pipeline_context_with_placement(
        &self,
        pipeline_group_id: PipelineGroupId,
        pipeline_id: PipelineId,
        core_id: usize,
        num_cores: usize,
        thread_id: usize,
        deployment_generation: u64,
        numa_node_id: usize,
    ) -> PipelineContext {
        PipelineContext::new_with_generation(
            self.clone(),
            PipelineContextParams {
                pipeline_group_id,
                pipeline_id,
                core_id,
                num_cores,
                thread_id,
                numa_node_id,
            },
            deployment_generation,
        )
    }

    /// Registers the engine-level entity for engine-wide metrics with no scope attributes.
    ///
    /// Returns the [`EntityKey`] to pass to
    /// [`EngineMetricsMonitor::new`](crate::engine_metrics::EngineMetricsMonitor::new).
    #[must_use]
    pub fn register_engine_entity(&self) -> EntityKey {
        self.telemetry_registry_handle
            .register_entity(EngineEntityAttributeSet)
    }

    /// Returns auto-detected OpenTelemetry resource semantic-convention attributes.
    #[must_use]
    pub fn resource_attributes(&self) -> Vec<(String, String)> {
        let mut attributes = Vec::new();
        if !self.host_id.is_empty() {
            attributes.push(("host.id".to_owned(), self.host_id.to_string()));
        }
        if !self.container_id.is_empty() {
            attributes.push(("container.id".to_owned(), self.container_id.to_string()));
        }
        if !self.process_instance_id.is_empty() {
            attributes.push((
                "service.instance.id".to_owned(),
                self.process_instance_id.to_string(),
            ));
        }
        attributes
    }

    /// Returns a handle to the telemetry registry.
    #[must_use]
    pub fn telemetry_registry(&self) -> TelemetryRegistryHandle {
        self.telemetry_registry_handle.clone()
    }

    /// Returns the shared process-wide memory pressure state.
    #[must_use]
    pub fn memory_pressure_state(&self) -> MemoryPressureState {
        self.memory_pressure_state.clone()
    }
}

impl From<&PipelineContextParams> for PipelineKey {
    fn from(params: &PipelineContextParams) -> Self {
        PipelineKey::new(params.pipeline_group_id.clone(), params.pipeline_id.clone())
    }
}

impl PipelineContext {
    /// Returns the engine [StateDirectory](crate::state_dir::StateDirectory),
    /// or `None` when no state root is configured.
    #[must_use]
    pub fn state_directory(&self) -> Option<&crate::state_dir::StateDirectory> {
        self.controller_context.state_directory()
    }

    /// Creates a new `PipelineContext`.
    #[allow(dead_code)]
    pub(crate) fn new(
        parent_ctx: ControllerContext,
        pipeline_context_params: PipelineContextParams,
    ) -> Self {
        Self::new_with_generation(parent_ctx, pipeline_context_params, 0)
    }

    /// Creates a new `PipelineContext` with an explicit deployment generation.
    pub(crate) fn new_with_generation(
        parent_ctx: ControllerContext,
        pipeline_context_params: PipelineContextParams,
        deployment_generation: u64,
    ) -> Self {
        Self {
            controller_context: parent_ctx,
            pipeline_context_params,
            deployment_generation,
            node_id: Default::default(),
            node_urn: Default::default(),
            node_kind: Default::default(),
            node_interests: Interests::empty(),
            node_duration_distribution: DistributionTier::Normal,
            node_telemetry_attrs: HashMap::new(),
            admission: crate::admission::AdmissionBinder::none(),
            pipeline_telemetry_attrs: HashMap::new(),
            internal_telemetry: None,
            node_names: Arc::new(HashMap::new()),
            topic_set: None,
            listener_group_snapshot: Arc::new(ListenerGroupSnapshot::empty()),
            compiled_context_bindings: Arc::new(CompiledContextBindings::empty()),
        }
    }

    /// Returns the pipeline group ID associated with this pipeline context.
    #[must_use]
    pub fn pipeline_group_id(&self) -> PipelineGroupId {
        self.pipeline_context_params.pipeline_group_id.clone()
    }

    /// Returns the pipeline ID associated with this pipeline context.
    #[must_use]
    pub fn pipeline_id(&self) -> PipelineId {
        self.pipeline_context_params.pipeline_id.clone()
    }

    /// Returns the pipeline key.
    #[must_use]
    pub fn pipeline_key(&self) -> PipelineKey {
        PipelineKey::from(&self.pipeline_context_params)
    }

    /// Returns the node ID.
    #[must_use]
    pub fn node_id(&self) -> ConfigNodeId {
        self.node_id.clone()
    }

    /// Returns the core ID associated with this pipeline context.
    #[must_use]
    pub const fn core_id(&self) -> usize {
        self.pipeline_context_params.core_id
    }

    /// Returns the deployment generation associated with this pipeline runtime.
    #[must_use]
    pub const fn deployment_generation(&self) -> u64 {
        self.deployment_generation
    }

    /// Returns the total number of cores allocated to this pipeline.
    ///
    /// This is useful for nodes that need to share resources (like disk budgets)
    /// across all cores running the same pipeline.
    #[must_use]
    pub const fn num_cores(&self) -> usize {
        self.pipeline_context_params.num_cores
    }

    /// Sets the internal telemetry settings for the Internal Telemetry Receiver.
    ///
    /// This is called by the pipeline factory when building the internal telemetry pipeline.
    /// The ITR factory will read these settings during node construction.
    pub fn set_internal_telemetry(&mut self, settings: InternalTelemetrySettings) {
        self.internal_telemetry = Some(settings);
    }

    /// Returns the internal telemetry settings, if configured.
    ///
    /// Only the Internal Telemetry Receiver factory uses this to obtain the logs
    /// channel and resource bytes it needs for operation.
    #[must_use]
    pub const fn internal_telemetry(&self) -> Option<&InternalTelemetrySettings> {
        self.internal_telemetry.as_ref()
    }

    /// Returns the shared process-wide memory pressure state.
    #[must_use]
    pub fn memory_pressure_state(&self) -> MemoryPressureState {
        self.controller_context.memory_pressure_state()
    }

    /// Returns the node's construction-time ingress admission binder.
    ///
    /// Participating components must bind exactly once during their factory
    /// `create` call. The returned local or shared gate is then retained for
    /// the component's runtime lifetime.
    #[must_use]
    pub const fn admission(&self) -> &crate::admission::AdmissionBinder {
        &self.admission
    }

    /// Attaches the node's resolved admission binding before factory construction.
    pub(crate) fn set_admission(&mut self, admission: crate::admission::AdmissionBinder) {
        self.admission = admission;
    }

    /// Sets the shared node-name-to-index mapping for this pipeline context.
    pub fn set_node_names(&mut self, node_names: NodeNameIndex) {
        self.node_names = node_names;
    }

    /// Sets the pipeline-scoped topic set resource.
    pub fn set_topic_set<T: Send + Sync + 'static>(
        &mut self,
        topic_set: crate::topic::TopicSet<T>,
    ) {
        self.topic_set = Some(Arc::new(topic_set));
    }

    /// Sets the controller-resolved listener-group placement snapshot.
    pub fn set_listener_group_snapshot(&mut self, snapshot: ListenerGroupSnapshot) {
        self.listener_group_snapshot = Arc::new(snapshot);
    }

    /// Sets a shared controller-resolved listener-group placement snapshot.
    pub fn set_listener_group_snapshot_arc(&mut self, snapshot: Arc<ListenerGroupSnapshot>) {
        self.listener_group_snapshot = snapshot;
    }

    /// Returns the controller-resolved listener-group placement snapshot.
    #[must_use]
    pub fn listener_group_snapshot(&self) -> Arc<ListenerGroupSnapshot> {
        Arc::clone(&self.listener_group_snapshot)
    }

    /// Sets this context's compiled bindings.
    pub fn set_compiled_context_bindings(&mut self, bindings: Arc<CompiledContextBindings>) {
        self.compiled_context_bindings = bindings;
    }

    /// Returns this context's compiled bindings.
    #[must_use]
    pub fn compiled_context_bindings(&self) -> &Arc<CompiledContextBindings> {
        &self.compiled_context_bindings
    }

    /// Returns the pipeline-scoped topic set, if one was injected.
    #[must_use]
    pub fn topic_set<T: Send + Sync + 'static>(&self) -> Option<crate::topic::TopicSet<T>> {
        self.topic_set
            .as_ref()
            .and_then(|resource| resource.downcast_ref::<crate::topic::TopicSet<T>>())
            .cloned()
    }

    /// Returns the pipeline index for the given node name, if it exists.
    #[must_use]
    pub fn node_by_name(&self, name: &str) -> Option<EngineNodeId> {
        self.node_names.get(name).cloned()
    }

    /// Takes the internal telemetry settings, leaving None in its place.
    ///
    /// Used by the ITR factory to consume the settings during construction.
    #[must_use]
    pub const fn take_internal_telemetry(&mut self) -> Option<InternalTelemetrySettings> {
        self.internal_telemetry.take()
    }

    /// Registers a metric set for the given entity key and tracks it in node telemetry if present.
    #[must_use]
    pub fn register_metric_set_for_entity<T: MetricSetHandler + Default + Debug + Send + Sync>(
        &self,
        entity_key: EntityKey,
    ) -> MetricSet<T> {
        let metrics = self
            .controller_context
            .telemetry_registry_handle
            .register_metric_set_for_entity::<T>(entity_key);
        if let Some(telemetry) = current_node_telemetry_handle() {
            telemetry.track_metric_set(metrics.metric_set_key());
        }
        metrics
    }

    /// Registers a measurement metric set for the given entity key and tracks it in node telemetry if present.
    #[must_use]
    pub fn register_measurement_metric_set_for_entity<
        T: MeasurementMetricSetHandler + Debug + Send + Sync,
    >(
        &self,
        entity_key: EntityKey,
    ) -> MeasurementMetricSet<T> {
        let metrics = self
            .controller_context
            .telemetry_registry_handle
            .register_metric_set_with_measurement_attributes_for_entity::<T>(entity_key);
        if let Some(telemetry) = current_node_telemetry_handle() {
            telemetry.track_metric_set(metrics.metric_set_key());
        }
        metrics
    }

    /// Returns a registrar for metric sets scoped to an existing entity.
    #[must_use]
    pub const fn metric_set_registrar_for_entity(
        &self,
        entity_key: EntityKey,
    ) -> EntityMetricSetRegistrar<'_> {
        EntityMetricSetRegistrar {
            pipeline_context: self,
            entity_key,
        }
    }

    /// Returns a registrar for metric sets scoped by an additional `topic` attribute.
    #[must_use]
    pub fn metric_set_registrar_with_topic(
        &self,
        topic: Cow<'static, str>,
    ) -> EntityMetricSetRegistrar<'_> {
        let entity_key = self.register_topic_entity(topic);
        if let Some(telemetry) = current_node_telemetry_handle() {
            telemetry.track_entity(entity_key);
        }
        self.metric_set_registrar_for_entity(entity_key)
    }

    /// Registers an entity and tracks it for cleanup with the current node, if present.
    #[must_use]
    pub fn register_entity(
        &self,
        attributes: impl AttributeSetHandler + Send + Sync + 'static,
    ) -> EntityKey {
        let entity_key = self
            .controller_context
            .telemetry_registry_handle
            .register_entity(attributes);
        if let Some(telemetry) = current_node_telemetry_handle() {
            telemetry.track_entity(entity_key);
        }
        entity_key
    }

    /// Shared entity-resolution skeleton for the `register_*_metrics` family.
    ///
    /// Resolves the current node's telemetry scope in priority order -- active node
    /// telemetry handle, then ambient node entity key -- and registers the metric set
    /// via `for_entity`. When registering against an active node handle, the resulting
    /// metric-set key (obtained through `metric_set_key`) is tracked so the set is
    /// unregistered as part of node cleanup.
    ///
    /// Tests often construct nodes directly without the engine's entity scoping, so a
    /// final fallback registers the set against this context's own attribute set via
    /// `with_scope` (test builds only); production builds panic.
    fn register_scoped_metrics<R>(
        &self,
        for_entity: impl FnOnce(&TelemetryRegistryHandle, EntityKey) -> R,
        metric_set_key: impl FnOnce(&R) -> MetricSetKey,
        with_scope: impl FnOnce(&Self, &TelemetryRegistryHandle) -> R,
    ) -> R {
        let handle = &self.controller_context.telemetry_registry_handle;
        if let Some(telemetry) = current_node_telemetry_handle() {
            let metrics = for_entity(handle, telemetry.entity_key());
            telemetry.track_metric_set(metric_set_key(&metrics));
            metrics
        } else if let Some(entity_key) = node_entity_key() {
            for_entity(handle, entity_key)
        } else {
            #[cfg(any(test, feature = "test-utils"))]
            {
                with_scope(self, handle)
            }

            #[cfg(not(any(test, feature = "test-utils")))]
            {
                let _ = with_scope;
                panic!(
                    "node entity key not set; ensure node entity is registered and instrumented"
                );
            }
        }
    }

    fn register_topic_entity(&self, topic: Cow<'static, str>) -> EntityKey {
        if self.node_telemetry_attrs.is_empty() {
            self.controller_context
                .telemetry_registry_handle
                .register_entity(NodeWithTopicAttributeSet {
                    node_attrs: self.node_attribute_set(),
                    topic,
                })
        } else {
            self.controller_context
                .telemetry_registry_handle
                .register_entity(NodeWithCustomTopicAttributeSet {
                    node_custom_attrs: self.node_with_custom_attribute_set(),
                    topic,
                })
        }
    }

    /// Registers the pipeline entity for this context.
    #[must_use]
    pub fn register_pipeline_entity(&self) -> EntityKey {
        self.controller_context
            .telemetry_registry_handle
            .register_entity(self.pipeline_attribute_set())
    }

    /// Returns an [`ExtensionContext`] scoped to this pipeline.
    #[must_use]
    pub fn extension_context(&self) -> ExtensionContext {
        ExtensionContext::new(
            self.controller_context.clone(),
            ExtensionScopeAttributeSet::pipeline(self.pipeline_attribute_set()),
        )
    }

    /// Registers the node entity for this context.
    ///
    /// If the node has custom telemetry attributes configured, they are included
    /// in the entity registration. Otherwise, only the base node attributes are used,
    /// keeping telemetry output clean for nodes without custom attributes.
    #[must_use]
    pub fn register_node_entity(&self) -> EntityKey {
        if self.node_telemetry_attrs.is_empty() {
            self.controller_context
                .telemetry_registry_handle
                .register_entity(self.node_attribute_set())
        } else {
            self.controller_context
                .telemetry_registry_handle
                .register_entity(self.node_with_custom_attribute_set())
        }
    }

    fn engine_attribute_set(&self) -> EngineAttributeSet {
        EngineAttributeSet {
            core_id: self.pipeline_context_params.core_id,
            numa_node_id: self.pipeline_context_params.numa_node_id,
        }
    }

    /// Returns the pipeline attribute set for the current pipeline context.
    #[must_use]
    pub fn pipeline_attribute_set(&self) -> PipelineAttributeSet {
        PipelineAttributeSet {
            engine_attrs: self.engine_attribute_set(),
            pipeline_id: self.pipeline_context_params.pipeline_id.clone(),
            pipeline_group_id: self.pipeline_context_params.pipeline_group_id.clone(),
            deployment_generation: self.deployment_generation,
        }
    }

    /// Returns the node attribute set for the current node context.
    #[must_use]
    pub fn node_attribute_set(&self) -> NodeAttributeSet {
        NodeAttributeSet {
            pipeline_attrs: self.pipeline_attribute_set(),
            node_id: self.node_id.clone(),
            node_urn: self.node_urn.clone().into(),
            node_type: self.node_kind.into(),
        }
    }

    /// Returns whether the node has custom telemetry identity attributes.
    #[must_use]
    pub fn has_custom_node_attributes(&self) -> bool {
        !self.node_telemetry_attrs.is_empty()
    }

    /// Returns the node attribute set extended with custom telemetry attributes.
    ///
    /// Only used when the node has non-empty `entity.extend.identity_attributes` configured.
    #[must_use]
    pub fn node_with_custom_attribute_set(&self) -> NodeWithCustomAttributeSet {
        NodeWithCustomAttributeSet {
            node_attrs: self.node_attribute_set(),
            custom_attrs: CustomAttributeSet::new(config_map_to_telemetry(
                &self.node_telemetry_attrs,
            )),
        }
    }

    /// Returns a channel attribute set tied to this node context.
    #[must_use]
    pub fn node_channel_attribute_set(
        &self,
        channel_id: Cow<'static, str>,
        node_port: Cow<'static, str>,
        channel_kind: ChannelKind,
        channel_mode: ChannelMode,
        channel_type: ChannelType,
        channel_impl: ChannelImplementation,
    ) -> NodeChannelAttributeSet {
        NodeChannelAttributeSet {
            node_attrs: self.node_attribute_set(),
            node_port,
            channel_id,
            channel_kind,
            channel_mode,
            channel_type,
            channel_impl,
        }
    }

    /// Registers a node-scoped channel entity for the given channel attributes.
    #[must_use]
    pub fn register_node_channel_entity(
        &self,
        channel_id: Cow<'static, str>,
        node_port: Cow<'static, str>,
        channel_kind: ChannelKind,
        channel_mode: ChannelMode,
        channel_type: ChannelType,
        channel_impl: ChannelImplementation,
    ) -> EntityKey {
        let attrs = self.node_channel_attribute_set(
            channel_id,
            node_port,
            channel_kind,
            channel_mode,
            channel_type,
            channel_impl,
        );
        let registry = &self.controller_context.telemetry_registry_handle;

        if self.node_telemetry_attrs.is_empty() {
            registry.register_entity(attrs)
        } else {
            registry.register_entity(NodeWithCustomChannelAttributeSet {
                channel_attrs: attrs,
                custom_attrs: CustomAttributeSet::new(config_map_to_telemetry(
                    &self.node_telemetry_attrs,
                )),
            })
        }
    }

    /// Returns a metrics registry handle.
    #[must_use]
    pub fn metrics_registry(&self) -> TelemetryRegistryHandle {
        self.controller_context.telemetry_registry_handle.clone()
    }

    /// Returns the effective telemetry interests for this node.
    #[must_use]
    pub const fn node_interests(&self) -> Interests {
        self.node_interests
    }

    /// Sets the effective telemetry interests for this node.
    pub(crate) fn set_node_interests(&mut self, interests: Interests) {
        self.node_interests = interests;
    }

    /// Returns the aggregation fidelity for node-local duration measurements.
    #[must_use]
    pub const fn node_duration_distribution(&self) -> DistributionTier {
        self.node_duration_distribution
    }

    /// Sets the aggregation fidelity for node-local duration measurements.
    pub(crate) fn set_node_duration_distribution(&mut self, tier: DistributionTier) {
        self.node_duration_distribution = tier;
    }

    /// Returns a new pipeline context with the given node identifiers.
    #[must_use]
    pub fn with_node_context(
        &self,
        node_id: ConfigNodeId,
        node_urn: NodeUrn,
        node_kind: NodeKind,
        node_telemetry_attrs: HashMap<String, TelemetryAttribute>,
    ) -> Self {
        Self {
            controller_context: self.controller_context.clone(),
            pipeline_context_params: self.pipeline_context_params.clone(),
            deployment_generation: self.deployment_generation,
            pipeline_telemetry_attrs: self.pipeline_telemetry_attrs.clone(),
            node_id,
            node_urn,
            node_kind,
            node_interests: Interests::empty(),
            node_duration_distribution: DistributionTier::Normal,
            node_telemetry_attrs,
            admission: crate::admission::AdmissionBinder::none(),
            internal_telemetry: None,
            node_names: self.node_names.clone(),
            topic_set: self.topic_set.clone(),
            listener_group_snapshot: Arc::clone(&self.listener_group_snapshot),
            compiled_context_bindings: self.compiled_context_bindings.clone(),
        }
    }
}

impl MetricSetRegistrar for EntityMetricSetRegistrar<'_> {
    fn register_metric_set<M: MetricSetHandler + Default + Debug + Send + Sync>(
        &self,
    ) -> MetricSet<M> {
        self.pipeline_context
            .register_metric_set_for_entity(self.entity_key)
    }

    fn register_registration_metric_set<M: RegistrationMetricSetHandler + Debug + Send + Sync>(
        &self,
        registration_attrs: &M::RegistrationAttributes,
    ) -> MetricSet<M> {
        let metrics = self
            .pipeline_context
            .controller_context
            .telemetry_registry_handle
            .register_metric_set_with_registration_attributes_for_entity::<M>(
                self.entity_key,
                registration_attrs,
            );
        if let Some(telemetry) = current_node_telemetry_handle() {
            telemetry.track_metric_set(metrics.metric_set_key());
        }
        metrics
    }

    fn register_measurement_metric_set<M: MeasurementMetricSetHandler + Debug + Send + Sync>(
        &self,
    ) -> MeasurementMetricSet<M> {
        self.pipeline_context
            .register_measurement_metric_set_for_entity(self.entity_key)
    }

    fn register_registration_and_measurement_metric_set<
        M: RegistrationMetricSetHandler + MeasurementMetricSetHandler + Debug + Send + Sync,
    >(
        &self,
        registration_attrs: &M::RegistrationAttributes,
    ) -> MeasurementMetricSet<M> {
        let metrics = self
            .pipeline_context
            .controller_context
            .telemetry_registry_handle
            .register_metric_set_with_registration_and_measurement_attributes_for_entity::<M>(
                self.entity_key,
                registration_attrs,
            );
        if let Some(telemetry) = current_node_telemetry_handle() {
            telemetry.track_metric_set(metrics.metric_set_key());
        }
        metrics
    }
}

impl MetricSetRegistrar for PipelineContext {
    fn register_metric_set<M: MetricSetHandler + Default + Debug + Send + Sync>(
        &self,
    ) -> MetricSet<M> {
        self.register_scoped_metrics(
            |handle, entity_key| handle.register_metric_set_for_entity::<M>(entity_key),
            MetricSet::metric_set_key,
            |ctx, handle| {
                if ctx.node_telemetry_attrs.is_empty() {
                    handle.register_metric_set::<M>(ctx.node_attribute_set())
                } else {
                    handle.register_metric_set::<M>(ctx.node_with_custom_attribute_set())
                }
            },
        )
    }

    fn register_registration_metric_set<M: RegistrationMetricSetHandler + Debug + Send + Sync>(
        &self,
        registration_attrs: &M::RegistrationAttributes,
    ) -> MetricSet<M> {
        self.register_scoped_metrics(
            |handle, entity_key| {
                handle.register_metric_set_with_registration_attributes_for_entity::<M>(
                    entity_key,
                    registration_attrs,
                )
            },
            MetricSet::metric_set_key,
            |ctx, handle| {
                if ctx.node_telemetry_attrs.is_empty() {
                    handle.register_metric_set_with_registration_attributes::<M>(
                        ctx.node_attribute_set(),
                        registration_attrs,
                    )
                } else {
                    handle.register_metric_set_with_registration_attributes::<M>(
                        ctx.node_with_custom_attribute_set(),
                        registration_attrs,
                    )
                }
            },
        )
    }

    fn register_measurement_metric_set<M: MeasurementMetricSetHandler + Debug + Send + Sync>(
        &self,
    ) -> MeasurementMetricSet<M> {
        self.register_scoped_metrics(
            |handle, entity_key| {
                handle.register_metric_set_with_measurement_attributes_for_entity::<M>(entity_key)
            },
            MeasurementMetricSet::metric_set_key,
            |ctx, handle| {
                if ctx.node_telemetry_attrs.is_empty() {
                    handle.register_metric_set_with_measurement_attributes::<M>(
                        ctx.node_attribute_set(),
                    )
                } else {
                    handle.register_metric_set_with_measurement_attributes::<M>(
                        ctx.node_with_custom_attribute_set(),
                    )
                }
            },
        )
    }

    fn register_registration_and_measurement_metric_set<
        M: RegistrationMetricSetHandler + MeasurementMetricSetHandler + Debug + Send + Sync,
    >(
        &self,
        registration_attrs: &M::RegistrationAttributes,
    ) -> MeasurementMetricSet<M> {
        self.register_scoped_metrics(
            |handle, entity_key| {
                handle
                    .register_metric_set_with_registration_and_measurement_attributes_for_entity::<M>(
                    entity_key,
                    registration_attrs,
                )
            },
            MeasurementMetricSet::metric_set_key,
            |ctx, handle| {
                if ctx.node_telemetry_attrs.is_empty() {
                    handle.register_metric_set_with_registration_and_measurement_attributes::<M>(
                        ctx.node_attribute_set(),
                        registration_attrs,
                    )
                } else {
                    handle.register_metric_set_with_registration_and_measurement_attributes::<M>(
                        ctx.node_with_custom_attribute_set(),
                        registration_attrs,
                    )
                }
            },
        )
    }
}

/// Host-scope context for extensions.
#[derive(Clone, Debug)]
pub struct ExtensionContext {
    controller_context: ControllerContext,
    extension_scope: ExtensionScopeAttributeSet,
}

impl ExtensionContext {
    /// Creates an `ExtensionContext` for the supplied host scope.
    ///
    /// # Panics (debug builds)
    ///
    /// Panics if `extension_scope.kind` is empty.
    #[must_use]
    pub fn new(
        controller_context: ControllerContext,
        extension_scope: ExtensionScopeAttributeSet,
    ) -> Self {
        debug_assert!(
            !extension_scope.kind.is_empty(),
            "ExtensionContext requires a non-empty scope kind"
        );
        Self {
            controller_context,
            extension_scope,
        }
    }

    /// Returns the telemetry registry handle for this scope.
    #[must_use]
    pub fn metrics_registry(&self) -> TelemetryRegistryHandle {
        self.controller_context.telemetry_registry_handle.clone()
    }

    /// Returns the attribute set for an extension hosted at this scope.
    #[must_use]
    pub fn extension_attribute_set(
        &self,
        extension_id: Cow<'static, str>,
        variant: crate::extension::wrapper::ExtensionVariant,
    ) -> ExtensionAttributeSet {
        ExtensionAttributeSet {
            extension_id,
            extension_variant: Cow::Borrowed(variant.as_str()),
            extension_scope: self.extension_scope.clone(),
        }
    }

    /// Registers an extension entity at this scope.
    #[must_use]
    pub fn register_extension_entity(
        &self,
        extension_id: Cow<'static, str>,
        variant: crate::extension::wrapper::ExtensionVariant,
    ) -> EntityKey {
        self.controller_context
            .telemetry_registry_handle
            .register_entity(self.extension_attribute_set(extension_id, variant))
    }

    /// Returns a channel attribute set tied to the given extension.
    #[must_use]
    pub fn extension_channel_attribute_set(
        &self,
        extension_id: Cow<'static, str>,
        variant: crate::extension::wrapper::ExtensionVariant,
        channel_id: Cow<'static, str>,
        channel_mode: ChannelMode,
        channel_impl: ChannelImplementation,
    ) -> ExtensionChannelAttributeSet {
        ExtensionChannelAttributeSet {
            extension_attrs: self.extension_attribute_set(extension_id, variant),
            channel_id,
            channel_mode,
            channel_impl,
        }
    }

    /// Registers an extension-scoped channel entity for the given attributes.
    #[must_use]
    pub fn register_extension_channel_entity(
        &self,
        extension_id: Cow<'static, str>,
        variant: crate::extension::wrapper::ExtensionVariant,
        channel_id: Cow<'static, str>,
        channel_mode: ChannelMode,
        channel_impl: ChannelImplementation,
    ) -> EntityKey {
        let attrs = self.extension_channel_attribute_set(
            extension_id,
            variant,
            channel_id,
            channel_mode,
            channel_impl,
        );
        self.controller_context
            .telemetry_registry_handle
            .register_entity(attrs)
    }

    /// Registers a metric set for the given entity key.
    ///
    /// Unlike [`PipelineContext::register_metric_set_for_entity`], this does
    /// not hook into any ambient node telemetry -- extension entities own their
    /// own lifecycle via the per-variant `EntityTelemetryGuard`.
    #[must_use]
    pub fn register_metric_set_for_entity<T: MetricSetHandler + Default + Debug + Send + Sync>(
        &self,
        entity_key: EntityKey,
    ) -> MetricSet<T> {
        self.controller_context
            .telemetry_registry_handle
            .register_metric_set_for_entity::<T>(entity_key)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use otel_arrow_dfe_config::pipeline::telemetry::AttributeValue;
    use otel_arrow_dfe_telemetry::registry::TelemetryRegistryHandle;
    use std::collections::HashMap;

    /// Scenario: contexts across cores and generations have no configured root.
    /// Guarantees: contexts have no state capability or fallback directory.
    #[test]
    fn state_directory_is_optional_across_generations() {
        let controller = ControllerContext::new(TelemetryRegistryHandle::new());
        assert!(controller.state_directory().is_none());
        for generation in [0, 1, 2] {
            let context = controller.pipeline_context_with_generation(
                "g".into(),
                "p".into(),
                generation as usize,
                3,
                generation as usize,
                generation,
            );
            assert!(context.state_directory().is_none());
        }
    }

    /// Scenario: explicit process, host, and container identities are available.
    /// Guarantees: resource attributes map all identities to stable semantic-convention keys.
    #[test]
    fn resource_attributes_maps_semconv_keys() {
        let ctx = ControllerContext::new_with_identity(
            TelemetryRegistryHandle::new(),
            "proc-123",
            "machine-abc",
            "container-xyz",
        );

        // All three identities present: emitted in a stable order under their
        // OpenTelemetry semantic-convention keys.
        assert_eq!(
            ctx.resource_attributes(),
            vec![
                ("host.id".to_string(), "machine-abc".to_string()),
                ("container.id".to_string(), "container-xyz".to_string()),
                ("service.instance.id".to_string(), "proc-123".to_string()),
            ]
        );
    }

    /// Scenario: host and container identity values are empty.
    /// Guarantees: resource attributes omit empty values while retaining the process identity.
    #[test]
    fn resource_attributes_omits_empty_values() {
        // Empty host/container should be skipped entirely (no empty-valued
        // attributes), leaving only the populated service.instance.id.
        let ctx = ControllerContext::new_with_identity(
            TelemetryRegistryHandle::new(),
            "proc-123",
            "",
            "",
        );

        assert_eq!(
            ctx.resource_attributes(),
            vec![("service.instance.id".to_string(), "proc-123".to_string())]
        );
    }

    /// Scenario: a pipeline context is created with an explicitly resolved NUMA node.
    /// Guarantees: engine telemetry reports both the selected core and NUMA node identifiers.
    #[test]
    fn pipeline_context_uses_resolved_numa_node_for_engine_attrs() {
        let ctx = ControllerContext::new(TelemetryRegistryHandle::new())
            .pipeline_context_with_placement(
                PipelineGroupId::from("g"),
                PipelineId::from("p"),
                4,
                1,
                0,
                0,
                2,
            );

        let attrs = ctx.pipeline_attribute_set();

        assert_eq!(attrs.engine_attrs.core_id, 4);
        assert_eq!(attrs.engine_attrs.numa_node_id, 2);
    }

    /// Scenario: a pipeline context with a listener plan is converted to a node context.
    /// Guarantees: the node context retains the listener snapshot generation and plan lookup.
    #[test]
    fn pipeline_context_preserves_listener_group_snapshot_across_node_context() {
        let addr = "127.0.0.1:4317".parse().unwrap();
        let mut ctx = ControllerContext::new(TelemetryRegistryHandle::new()).pipeline_context_with(
            Default::default(),
            Default::default(),
            0,
            1,
            0,
        );
        ctx.set_listener_group_snapshot(ListenerGroupSnapshot::new(
            3,
            vec![crate::listener_group::ListenerGroupPlan {
                key: crate::listener_group::ListenerGroupKey::new(
                    "group".into(),
                    "pipeline".into(),
                    "receiver".into(),
                    addr,
                    crate::listener_group::ListenerProtocol::Tcp,
                ),
                expected_members: Vec::new(),
            }],
        ));

        let node_ctx = ctx.with_node_context(
            "receiver".into(),
            "urn:otel:receiver:otlp".into(),
            NodeKind::Receiver,
            HashMap::new(),
        );

        assert_eq!(node_ctx.listener_group_snapshot().generation, 3);
        assert!(
            node_ctx
                .listener_group_snapshot()
                .plan_for(
                    "receiver",
                    addr,
                    crate::listener_group::ListenerProtocol::Tcp
                )
                .is_some()
        );
    }

    /// Scenario: a node context is created from a pipeline context.
    /// Guarantees: both contexts share the same compiled binding snapshot.
    #[test]
    fn pipeline_context_preserves_compiled_policy_across_node_context() {
        let resolved = otel_arrow_dfe_config::engine::ResolvedOtelDataflowSpec {
            engine: Default::default(),
            pipelines: Vec::new(),
        };
        let factory = crate::PipelineFactory::<()>::new(&[], &[], &[], &[]);
        let bindings = factory
            .compile_initial_context(&resolved)
            .expect("context bindings")
            .bindings;
        let controller = ControllerContext::new(TelemetryRegistryHandle::new());
        let mut pipeline =
            controller.pipeline_context_with("group".into(), "pipeline".into(), 0, 1, 0);
        pipeline.set_compiled_context_bindings(Arc::clone(&bindings));

        let node = pipeline.with_node_context(
            "node".into(),
            "urn:otel:processor:test".into(),
            NodeKind::Processor,
            HashMap::new(),
        );

        assert!(Arc::ptr_eq(node.compiled_context_bindings(), &bindings));
    }

    fn pipeline_ctx_with_custom_attrs(
        registry: TelemetryRegistryHandle,
        custom: HashMap<String, TelemetryAttribute>,
    ) -> PipelineContext {
        let controller_ctx = ControllerContext::new(registry);
        let pipeline_params = PipelineContextParams {
            pipeline_group_id: Cow::Borrowed("group1"),
            pipeline_id: Cow::Borrowed("pipe1"),
            core_id: 0,
            num_cores: 1,
            thread_id: 0,
            numa_node_id: 0,
        };
        PipelineContext::new(controller_ctx, pipeline_params).with_node_context(
            Cow::Borrowed("test-node"),
            NodeUrn::parse("urn:otel:receiver:test").unwrap(),
            NodeKind::Receiver,
            custom,
        )
    }

    fn register_channel(ctx: &PipelineContext) -> EntityKey {
        ctx.register_node_channel_entity(
            Cow::Borrowed("channel-1"),
            Cow::Borrowed("out"),
            ChannelKind::Pdata,
            ChannelMode::Local,
            ChannelType::Mpsc,
            ChannelImplementation::Internal,
        )
    }

    /// Scenario: a node configured with `entity.extend.identity_attributes` registers a
    /// channel endpoint entity.
    /// Guarantees: the channel entity carries the configured custom attributes in addition to
    /// the base channel attributes, so channel metrics keep the node's configured identity.
    #[test]
    fn register_node_channel_entity_includes_custom_attributes() {
        let registry = TelemetryRegistryHandle::new();
        let mut custom = HashMap::new();
        let _ = custom.insert(
            "custom.identity.foo".to_string(),
            TelemetryAttribute::new(AttributeValue::String("bar".to_string())),
        );
        let ctx = pipeline_ctx_with_custom_attrs(registry.clone(), custom);
        let key = register_channel(&ctx);

        let (schema, rendered) = registry
            .visit_entity(key, |a| (a.schema_name(), a.attributes_to_string()))
            .expect("channel entity registered");

        assert_eq!(schema, "node.channel.custom.attrs");
        assert!(
            rendered.contains("custom={custom.identity.foo=bar}"),
            "custom identity attributes missing from channel entity: {rendered}"
        );
        assert!(
            rendered.contains("channel.id=channel-1") && rendered.contains("node.id=test-node"),
            "base channel attributes must be preserved: {rendered}"
        );
    }

    /// Scenario: a node with no `entity.extend.identity_attributes` registers a channel
    /// endpoint entity.
    /// Guarantees: the entity stays on the plain channel schema and emits no empty
    /// `custom={}` attribute, keeping telemetry output clean for unconfigured nodes.
    #[test]
    fn register_node_channel_entity_omits_empty_custom_attributes() {
        let registry = TelemetryRegistryHandle::new();
        let ctx = pipeline_ctx_with_custom_attrs(registry.clone(), HashMap::new());
        let key = register_channel(&ctx);

        let (schema, rendered) = registry
            .visit_entity(key, |a| (a.schema_name(), a.attributes_to_string()))
            .expect("channel entity registered");

        assert_eq!(schema, "node.channel.attrs");
        assert!(
            !rendered.contains("custom="),
            "nodes without custom attributes must not emit a custom attribute: {rendered}"
        );
    }

    fn register_test_topic(ctx: &PipelineContext) -> EntityKey {
        ctx.register_topic_entity(Cow::Borrowed("test-topic"))
    }

    /// Scenario: a node configured with `entity.extend.identity_attributes` registers a topic entity.
    /// Guarantees: the topic entity carries the configured custom attributes along with the topic name.
    #[test]
    fn register_topic_entity_includes_custom_attributes() {
        let registry = TelemetryRegistryHandle::new();
        let mut custom = HashMap::new();
        let _ = custom.insert(
            "custom.identity.foo".to_string(),
            TelemetryAttribute::new(AttributeValue::String("bar".to_string())),
        );
        let ctx = pipeline_ctx_with_custom_attrs(registry.clone(), custom);
        let key = register_test_topic(&ctx);

        let (schema, rendered) = registry
            .visit_entity(key, |a| (a.schema_name(), a.attributes_to_string()))
            .expect("topic entity registered");

        assert_eq!(schema, "node.custom.topic.attrs");
        assert!(
            rendered.contains("custom={custom.identity.foo=bar}"),
            "custom identity attributes missing from topic entity: {rendered}"
        );
        assert!(
            rendered.contains("topic=test-topic") && rendered.contains("node.id=test-node"),
            "base topic attributes must be preserved: {rendered}"
        );
    }

    /// Scenario: a node with no `entity.extend.identity_attributes` registers a topic entity.
    /// Guarantees: the entity stays on the plain topic schema and emits no empty
    /// `custom={}` attribute.
    #[test]
    fn register_topic_entity_omits_empty_custom_attributes() {
        let registry = TelemetryRegistryHandle::new();
        let ctx = pipeline_ctx_with_custom_attrs(registry.clone(), HashMap::new());
        let key = register_test_topic(&ctx);

        let (schema, rendered) = registry
            .visit_entity(key, |a| (a.schema_name(), a.attributes_to_string()))
            .expect("topic entity registered");

        assert_eq!(schema, "node.topic.attrs");
        assert!(
            !rendered.contains("custom="),
            "nodes without custom attributes must not emit a custom attribute: {rendered}"
        );
    }

    /// Scenario: a node registers a measurement metric set through a topic registrar.
    /// Guarantees: the registered measurement set links to the topic entity, preserving
    /// topic, node, and custom identity attributes when configured, and omitting empty custom attributes when not.
    #[test]
    fn topic_metric_registrar_links_entity_with_and_without_custom_attrs() {
        use crate::flow_metrics::FlowInputMessageMetrics;

        // Without custom attributes
        let registry = TelemetryRegistryHandle::new();
        let ctx = pipeline_ctx_with_custom_attrs(registry.clone(), HashMap::new());
        let registrar = ctx.metric_set_registrar_with_topic(Cow::Borrowed("test-topic"));
        let metrics = FlowInputMessageMetrics::register(&registrar);
        let key = metrics.entity_key();
        let (schema, rendered) = registry
            .visit_entity(key, |a| (a.schema_name(), a.attributes_to_string()))
            .expect("measurement set entity registered without custom attrs");
        assert_eq!(schema, "node.topic.attrs");
        assert!(
            rendered.contains("topic=test-topic") && rendered.contains("node.id=test-node"),
            "base topic attributes must be preserved: {rendered}"
        );
        assert!(
            !rendered.contains("custom="),
            "nodes without custom attributes must not emit a custom attribute: {rendered}"
        );

        // With custom attributes
        let registry = TelemetryRegistryHandle::new();
        let mut custom = HashMap::new();
        let _ = custom.insert(
            "custom.identity.foo".to_string(),
            TelemetryAttribute::new(AttributeValue::String("bar".to_string())),
        );
        let ctx = pipeline_ctx_with_custom_attrs(registry.clone(), custom);
        let registrar = ctx.metric_set_registrar_with_topic(Cow::Borrowed("test-topic"));
        let metrics = FlowInputMessageMetrics::register(&registrar);
        let key = metrics.entity_key();
        let (schema, rendered) = registry
            .visit_entity(key, |a| (a.schema_name(), a.attributes_to_string()))
            .expect("measurement set entity registered with custom attrs");
        assert_eq!(schema, "node.custom.topic.attrs");
        assert!(
            rendered.contains("custom={custom.identity.foo=bar}"),
            "custom identity attributes missing from topic entity: {rendered}"
        );
        assert!(
            rendered.contains("topic=test-topic") && rendered.contains("node.id=test-node"),
            "base topic attributes must be preserved: {rendered}"
        );
    }

    /// Scenario: a node registers an arbitrary child entity and metric set through generic APIs.
    /// Guarantees: node cleanup unregisters both the child entity and its entity-bound metric set.
    #[test]
    fn generic_entity_registration_tracks_node_cleanup() {
        use crate::entity_context::{
            NodeTelemetryGuard, NodeTelemetryHandle, with_node_telemetry_handle,
        };
        use crate::flow_metrics::FlowInputMessageMetrics;

        let registry = TelemetryRegistryHandle::new();
        let ctx = pipeline_ctx_with_custom_attrs(registry.clone(), HashMap::new());
        let node_entity = ctx.register_node_entity();
        let handle = NodeTelemetryHandle::new(registry.clone(), node_entity);
        let guard = NodeTelemetryGuard::new(handle.clone());

        with_node_telemetry_handle(handle, || {
            let child_entity = ctx.register_entity(NodeWithTopicAttributeSet {
                node_attrs: ctx.node_attribute_set(),
                topic: Cow::Borrowed("child"),
            });
            let registrar = ctx.metric_set_registrar_for_entity(child_entity);
            let _metrics = FlowInputMessageMetrics::register(&registrar);
        });

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

        drop(guard);

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