otel-arrow-dfe-engine 0.61.0

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

//! Exporter wrapper used to provide a unified interface to the pipeline engine that abstracts over
//! the fact that exporter implementations may be `!Send` or `Send`.
//!
//! For more details on the `!Send` implementation of an exporter, see [`local::Exporter`].
//! See [`shared::Exporter`] for the Send implementation.

use crate::Interests;
use crate::channel_metrics::ChannelMetricsRegistry;
use crate::channel_mode::{LocalMode, SharedMode, wrap_node_control_channel_metrics};
use crate::completion_emission_metrics::CompletionEmissionMetricsHandle;
use crate::config::ExporterConfig;
use crate::context::PipelineContext;
use crate::context_declaration::CompiledHeaderPropagationPolicy as HeaderPropagationPolicy;
use crate::control::{
    Controllable, NodeControlMsg, PipelineCompletionMsgSender, RuntimeCtrlMsgSender,
};
use crate::entity_context::NodeTelemetryGuard;
use crate::error::{Error, ExporterErrorKind};
use crate::local::exporter as local;
use crate::local::message::{LocalReceiver, LocalSender};
use crate::message::{ExporterInbox, Receiver, Sender};
use crate::node::{Node, NodeId, NodeWithPDataReceiver};
use crate::runtime_services::PipelineRuntimeServices;
use crate::shared::exporter as shared;
use crate::shared::message::{SharedReceiver, SharedSender};
use crate::terminal_state::TerminalState;
use otel_arrow_dfe_channel::error::SendError;
use otel_arrow_dfe_channel::mpsc;
use otel_arrow_dfe_config::node::NodeUserConfig;
use otel_arrow_dfe_telemetry::reporter::MetricsReporter;
use std::sync::Arc;

/// A wrapper for the exporter that allows for both `Send` and `!Send` effect handlers.
///
/// Note: This is useful for creating a single interface for the exporter regardless of their
/// 'sendability'.
pub enum ExporterWrapper<PData> {
    /// An exporter with a `!Send` implementation.
    Local {
        /// Index identifier for the node.
        node_id: NodeId,
        /// The user configuration for the node, including its name and channel settings.
        user_config: Arc<NodeUserConfig>,
        /// The runtime configuration for the exporter.
        runtime_config: ExporterConfig,
        /// The exporter instance.
        exporter: Box<dyn local::Exporter<PData>>,
        /// A sender for control messages.
        control_sender: LocalSender<NodeControlMsg<PData>>,
        /// A receiver for control messages.
        control_receiver: LocalReceiver<NodeControlMsg<PData>>,
        /// Receiver for PData messages.
        pdata_receiver: Option<Receiver<PData>>,
        /// Telemetry guard for node lifecycle cleanup.
        telemetry: Option<NodeTelemetryGuard>,
        /// Pre-resolved propagation policy for transport header forwarding.
        propagation_policy: Option<HeaderPropagationPolicy>,
    },
    /// An exporter with a `Send` implementation.
    Shared {
        /// Index identifier for the node.
        node_id: NodeId,
        /// The user configuration for the node, including its name and channel settings.
        user_config: Arc<NodeUserConfig>,
        /// The runtime configuration for the exporter.
        runtime_config: ExporterConfig,
        /// The exporter instance.
        exporter: Box<dyn shared::Exporter<PData>>,
        /// A sender for control messages.
        control_sender: SharedSender<NodeControlMsg<PData>>,
        /// A receiver for control messages.
        control_receiver: SharedReceiver<NodeControlMsg<PData>>,
        /// Receiver for PData messages.
        pdata_receiver: Option<SharedReceiver<PData>>,
        /// Telemetry guard for node lifecycle cleanup.
        telemetry: Option<NodeTelemetryGuard>,
        /// Pre-resolved propagation policy for transport header forwarding.
        propagation_policy: Option<HeaderPropagationPolicy>,
    },
}

#[async_trait::async_trait(?Send)]
impl<PData> Controllable<PData> for ExporterWrapper<PData> {
    /// Returns the control message sender for the exporter.
    fn control_sender(&self) -> Sender<NodeControlMsg<PData>> {
        match self {
            ExporterWrapper::Local { control_sender, .. } => Sender::Local(control_sender.clone()),
            ExporterWrapper::Shared { control_sender, .. } => {
                Sender::Shared(control_sender.clone())
            }
        }
    }
}

impl<PData> ExporterWrapper<PData> {
    /// Creates a new local `ExporterWrapper` with the given exporter and configuration (!Send
    /// implementation).
    pub fn local<E>(
        exporter: E,
        node_id: NodeId,
        user_config: Arc<NodeUserConfig>,
        config: &ExporterConfig,
    ) -> Self
    where
        E: local::Exporter<PData> + 'static,
    {
        let (control_sender, control_receiver) =
            mpsc::Channel::new(config.control_channel.capacity);

        ExporterWrapper::Local {
            node_id,
            user_config,
            runtime_config: config.clone(),
            exporter: Box::new(exporter),
            control_sender: LocalSender::mpsc(control_sender),
            control_receiver: LocalReceiver::mpsc(control_receiver),
            pdata_receiver: None, // This will be set later
            telemetry: None,
            propagation_policy: None,
        }
    }

    /// Creates a new shared `ExporterWrapper` with the given exporter and configuration (Send
    /// implementation).
    pub fn shared<E>(
        exporter: E,
        node_id: NodeId,
        user_config: Arc<NodeUserConfig>,
        config: &ExporterConfig,
    ) -> Self
    where
        E: shared::Exporter<PData> + 'static,
    {
        let (control_sender, control_receiver) =
            tokio::sync::mpsc::channel(config.control_channel.capacity);

        ExporterWrapper::Shared {
            node_id,
            user_config,
            runtime_config: config.clone(),
            exporter: Box::new(exporter),
            control_sender: SharedSender::mpsc(control_sender),
            control_receiver: SharedReceiver::mpsc(control_receiver),
            pdata_receiver: None, // This will be set later
            telemetry: None,
            propagation_policy: None,
        }
    }

    pub(crate) fn with_node_telemetry_guard(self, guard: NodeTelemetryGuard) -> Self {
        match self {
            ExporterWrapper::Local {
                node_id,
                user_config,
                runtime_config,
                exporter,
                control_sender,
                control_receiver,
                pdata_receiver,
                propagation_policy,
                ..
            } => ExporterWrapper::Local {
                node_id,
                user_config,
                runtime_config,
                exporter,
                control_sender,
                control_receiver,
                pdata_receiver,
                telemetry: Some(guard),
                propagation_policy,
            },
            ExporterWrapper::Shared {
                node_id,
                user_config,
                runtime_config,
                exporter,
                control_sender,
                control_receiver,
                pdata_receiver,
                propagation_policy,
                ..
            } => ExporterWrapper::Shared {
                node_id,
                user_config,
                runtime_config,
                exporter,
                control_sender,
                control_receiver,
                pdata_receiver,
                telemetry: Some(guard),
                propagation_policy,
            },
        }
    }

    pub(crate) const fn take_telemetry_guard(&mut self) -> Option<NodeTelemetryGuard> {
        match self {
            ExporterWrapper::Local { telemetry, .. } => telemetry.take(),
            ExporterWrapper::Shared { telemetry, .. } => telemetry.take(),
        }
    }

    pub(crate) fn with_control_channel_metrics(
        self,
        pipeline_ctx: &PipelineContext,
        channel_metrics: &mut ChannelMetricsRegistry,
        channel_metrics_enabled: bool,
    ) -> Self {
        match self {
            ExporterWrapper::Local {
                node_id,
                runtime_config,
                control_sender,
                control_receiver,
                user_config,
                exporter,
                pdata_receiver,
                telemetry,
                propagation_policy,
                ..
            } => {
                let (control_sender, control_receiver) =
                    wrap_node_control_channel_metrics::<LocalMode, NodeControlMsg<PData>>(
                        node_id.name.as_ref(),
                        pipeline_ctx,
                        channel_metrics,
                        channel_metrics_enabled,
                        runtime_config.control_channel.capacity as u64,
                        control_sender,
                        control_receiver,
                    );

                ExporterWrapper::Local {
                    node_id,
                    user_config,
                    runtime_config,
                    exporter,
                    control_sender,
                    control_receiver,
                    pdata_receiver,
                    telemetry,
                    propagation_policy,
                }
            }
            ExporterWrapper::Shared {
                node_id,
                runtime_config,
                control_sender,
                control_receiver,
                user_config,
                exporter,
                pdata_receiver,
                telemetry,
                propagation_policy,
                ..
            } => {
                let (control_sender, control_receiver) =
                    wrap_node_control_channel_metrics::<SharedMode, NodeControlMsg<PData>>(
                        node_id.name.as_ref(),
                        pipeline_ctx,
                        channel_metrics,
                        channel_metrics_enabled,
                        runtime_config.control_channel.capacity as u64,
                        control_sender,
                        control_receiver,
                    );

                ExporterWrapper::Shared {
                    node_id,
                    user_config,
                    runtime_config,
                    exporter,
                    control_sender,
                    control_receiver,
                    pdata_receiver,
                    telemetry,
                    propagation_policy,
                }
            }
        }
    }

    /// Starts the exporter using the services owned by its pipeline runtime.
    pub async fn start(
        self,
        runtime_ctrl_msg_tx: RuntimeCtrlMsgSender<PData>,
        pipeline_completion_msg_tx: PipelineCompletionMsgSender<PData>,
        metrics_reporter: MetricsReporter,
        node_interests: Interests,
        runtime_services: PipelineRuntimeServices,
    ) -> Result<TerminalState, Error> {
        self.start_with_completion_metrics(
            runtime_ctrl_msg_tx,
            pipeline_completion_msg_tx,
            metrics_reporter,
            node_interests,
            None,
            runtime_services,
        )
        .await
    }

    pub(crate) async fn start_with_completion_metrics(
        self,
        runtime_ctrl_msg_tx: RuntimeCtrlMsgSender<PData>,
        pipeline_completion_msg_tx: PipelineCompletionMsgSender<PData>,
        metrics_reporter: MetricsReporter,
        node_interests: Interests,
        completion_emission_metrics: Option<CompletionEmissionMetricsHandle>,
        runtime_services: PipelineRuntimeServices,
    ) -> Result<TerminalState, Error> {
        match (self, metrics_reporter) {
            (
                ExporterWrapper::Local {
                    node_id,
                    exporter,
                    control_receiver,
                    pdata_receiver,
                    propagation_policy,
                    ..
                },
                metrics_reporter,
            ) => {
                let mut effect_handler = local::EffectHandler::new(
                    node_id.clone(),
                    metrics_reporter,
                    runtime_services.clone(),
                );
                let pdata_rx = pdata_receiver.ok_or_else(|| Error::ExporterError {
                    exporter: effect_handler.exporter_id(),
                    kind: ExporterErrorKind::Configuration,
                    error: "PData receiver not initialized".to_owned(),
                    source_detail: String::new(),
                })?;
                effect_handler
                    .core
                    .set_runtime_ctrl_msg_sender(runtime_ctrl_msg_tx);
                effect_handler
                    .core
                    .set_pipeline_completion_msg_sender(pipeline_completion_msg_tx);
                effect_handler.core.set_node_interests(node_interests);
                effect_handler
                    .core
                    .set_completion_emission_metrics(completion_emission_metrics.clone());
                effect_handler.set_propagation_policy(propagation_policy);
                let inbox = ExporterInbox::new(
                    Receiver::Local(control_receiver),
                    pdata_rx,
                    node_id.index,
                    node_interests,
                );
                exporter.start(inbox, effect_handler).await
            }
            (
                ExporterWrapper::Shared {
                    node_id,
                    exporter,
                    control_receiver,
                    pdata_receiver,
                    propagation_policy,
                    ..
                },
                metrics_reporter,
            ) => {
                let mut effect_handler =
                    shared::EffectHandler::new(node_id.clone(), metrics_reporter, runtime_services);
                let pdata_rx = pdata_receiver.ok_or_else(|| Error::ExporterError {
                    exporter: effect_handler.exporter_id(),
                    kind: ExporterErrorKind::Configuration,
                    error: "PData receiver not initialized".to_owned(),
                    source_detail: String::new(),
                })?;
                effect_handler
                    .core
                    .set_runtime_ctrl_msg_sender(runtime_ctrl_msg_tx);
                effect_handler
                    .core
                    .set_pipeline_completion_msg_sender(pipeline_completion_msg_tx);
                effect_handler.core.set_node_interests(node_interests);
                effect_handler
                    .core
                    .set_completion_emission_metrics(completion_emission_metrics);
                effect_handler.set_propagation_policy(propagation_policy);
                let inbox = shared::ExporterInbox::new(
                    control_receiver,
                    pdata_rx,
                    node_id.index,
                    node_interests,
                );
                exporter.start(inbox, effect_handler).await
            }
        }
    }

    /// Returns the wrapper with the given pre-resolved propagation policy for
    /// transport header forwarding.
    pub(crate) fn with_propagation_policy(self, policy: Option<HeaderPropagationPolicy>) -> Self {
        match self {
            ExporterWrapper::Local {
                node_id,
                user_config,
                runtime_config,
                exporter,
                control_sender,
                control_receiver,
                pdata_receiver,
                telemetry,
                ..
            } => ExporterWrapper::Local {
                node_id,
                user_config,
                runtime_config,
                exporter,
                control_sender,
                control_receiver,
                pdata_receiver,
                telemetry,
                propagation_policy: policy,
            },
            ExporterWrapper::Shared {
                node_id,
                user_config,
                runtime_config,
                exporter,
                control_sender,
                control_receiver,
                pdata_receiver,
                telemetry,
                ..
            } => ExporterWrapper::Shared {
                node_id,
                user_config,
                runtime_config,
                exporter,
                control_sender,
                control_receiver,
                pdata_receiver,
                telemetry,
                propagation_policy: policy,
            },
        }
    }
}

#[async_trait::async_trait(?Send)]
impl<PData> Node<PData> for ExporterWrapper<PData> {
    fn is_shared(&self) -> bool {
        match self {
            ExporterWrapper::Local { .. } => false,
            ExporterWrapper::Shared { .. } => true,
        }
    }

    fn node_id(&self) -> NodeId {
        match self {
            ExporterWrapper::Local { node_id, .. } => node_id.clone(),
            ExporterWrapper::Shared { node_id, .. } => node_id.clone(),
        }
    }

    fn user_config(&self) -> Arc<NodeUserConfig> {
        match self {
            ExporterWrapper::Local {
                user_config: config,
                ..
            } => config.clone(),
            ExporterWrapper::Shared {
                user_config: config,
                ..
            } => config.clone(),
        }
    }

    /// Sends a control message to the node.
    async fn send_control_msg(
        &self,
        msg: NodeControlMsg<PData>,
    ) -> Result<(), SendError<NodeControlMsg<PData>>> {
        match self {
            ExporterWrapper::Local { control_sender, .. } => control_sender.send(msg).await,
            ExporterWrapper::Shared { control_sender, .. } => control_sender.send(msg).await,
        }
    }
}

impl<PData> NodeWithPDataReceiver<PData> for ExporterWrapper<PData> {
    fn set_pdata_receiver(
        &mut self,
        node_id: NodeId,
        receiver: Receiver<PData>,
    ) -> Result<(), Error> {
        match (self, receiver) {
            (ExporterWrapper::Local { pdata_receiver, .. }, receiver) => {
                *pdata_receiver = Some(receiver);
                Ok(())
            }
            (ExporterWrapper::Shared { pdata_receiver, .. }, Receiver::Shared(receiver)) => {
                *pdata_receiver = Some(receiver);
                Ok(())
            }
            (ExporterWrapper::Shared { .. }, _) => Err(Error::ExporterError {
                exporter: node_id,
                kind: ExporterErrorKind::Configuration,
                error: "Expected a shared receiver for PData".to_owned(),
                source_detail: String::new(),
            }),
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::Interests;
    use crate::control::{AckMsg, NodeControlMsg};
    use crate::error::ExporterErrorKind;
    use crate::exporter::{Error, ExporterWrapper};
    use crate::local::exporter as local;
    use crate::local::message::LocalReceiver;
    use crate::message::{ExporterInbox, Message, ProcessorInbox, Receiver};
    use crate::shared::exporter as shared;
    use crate::shared::message::SharedReceiver;
    use crate::terminal_state::TerminalState;
    use crate::testing::exporter::TestContext;
    use crate::testing::exporter::TestRuntime;
    use crate::testing::{CtrlMsgCounters, TestMsg, test_node};
    use async_trait::async_trait;
    use otel_arrow_dfe_channel::error::RecvError;
    use otel_arrow_dfe_channel::mpsc;
    use otel_arrow_dfe_config::node::NodeUserConfig;
    use serde_json::Value;
    use std::future::Future;
    use std::ops::Add;
    use std::sync::Arc;
    use std::time::{Duration, Instant};
    use tokio::time::sleep;

    /// A test exporter that counts message events.
    /// Works with any type of exporter !Send or Send.
    pub struct TestExporter {
        /// Counter for different message types
        pub counter: CtrlMsgCounters,
    }

    impl TestExporter {
        /// Creates a new test node with the given counter
        pub fn new(counter: CtrlMsgCounters) -> Self {
            TestExporter { counter }
        }
    }

    #[async_trait(?Send)]
    impl local::Exporter<TestMsg> for TestExporter {
        async fn start(
            self: Box<Self>,
            mut msg_chan: ExporterInbox<TestMsg>,
            effect_handler: local::EffectHandler<TestMsg>,
        ) -> Result<TerminalState, Error> {
            // Loop until a Shutdown event is received.
            loop {
                match msg_chan.recv().await? {
                    Message::Control(NodeControlMsg::TimerTick { .. }) => {
                        self.counter.increment_timer_tick();
                    }
                    Message::Control(NodeControlMsg::Config { .. }) => {
                        self.counter.increment_config();
                    }
                    Message::Control(NodeControlMsg::Shutdown { .. }) => {
                        self.counter.increment_shutdown();
                        break;
                    }
                    Message::PData(_message) => {
                        self.counter.increment_message();
                    }
                    _ => {
                        return Err(Error::ExporterError {
                            exporter: effect_handler.exporter_id(),
                            kind: ExporterErrorKind::Other,
                            error: "Unknown control message".to_owned(),
                            source_detail: String::new(),
                        });
                    }
                }
            }
            Ok(TerminalState::default())
        }
    }

    #[async_trait]
    impl shared::Exporter<TestMsg> for TestExporter {
        async fn start(
            self: Box<Self>,
            mut msg_chan: shared::ExporterInbox<TestMsg>,
            effect_handler: shared::EffectHandler<TestMsg>,
        ) -> Result<TerminalState, Error> {
            // Loop until a Shutdown event is received.
            loop {
                match msg_chan.recv().await? {
                    Message::Control(NodeControlMsg::TimerTick { .. }) => {
                        self.counter.increment_timer_tick();
                    }
                    Message::Control(NodeControlMsg::Config { .. }) => {
                        self.counter.increment_config();
                    }
                    Message::Control(NodeControlMsg::Shutdown { .. }) => {
                        self.counter.increment_shutdown();
                        break;
                    }
                    Message::PData(_message) => {
                        self.counter.increment_message();
                    }
                    _ => {
                        return Err(Error::ExporterError {
                            exporter: effect_handler.exporter_id(),
                            kind: ExporterErrorKind::Other,
                            error: "Unknown control message".to_owned(),
                            source_detail: String::new(),
                        });
                    }
                }
            }
            Ok(TerminalState::default())
        }
    }

    /// Test closure that simulates a typical test scenario by sending timer ticks, config,
    /// data message, and shutdown control messages.
    fn scenario() -> impl FnOnce(TestContext<TestMsg>) -> std::pin::Pin<Box<dyn Future<Output = ()>>>
    {
        |ctx| {
            Box::pin(async move {
                // Send 3 TimerTick events.
                for _ in 0..3 {
                    ctx.send_timer_tick()
                        .await
                        .expect("Failed to send TimerTick");
                    ctx.sleep(Duration::from_millis(50)).await;
                }

                // Send a Config event.
                ctx.send_config(Value::Null)
                    .await
                    .expect("Failed to send Config");

                // Send a data message
                ctx.send_pdata(TestMsg("Hello Exporter".into()))
                    .await
                    .expect("Failed to send data message");

                // Send shutdown
                ctx.send_shutdown(
                    Instant::now().add(Duration::from_millis(200)),
                    "test complete",
                )
                .await
                .expect("Failed to send Shutdown");
            })
        }
    }

    /// Validation closure that checks the expected counter values
    fn validation_procedure() -> impl FnOnce(
        TestContext<TestMsg>,
        Result<(), Error>,
    ) -> std::pin::Pin<Box<dyn Future<Output = ()>>> {
        |ctx, _| {
            Box::pin(async move {
                ctx.counters().assert(
                    3, // timer tick
                    1, // message
                    1, // config
                    1, // shutdown
                );
            })
        }
    }

    #[test]
    fn test_exporter_local() {
        let test_runtime = TestRuntime::new();
        let user_config = Arc::new(NodeUserConfig::new_exporter_config("test_exporter"));
        let exporter = ExporterWrapper::local(
            TestExporter::new(test_runtime.counters()),
            test_node(test_runtime.config().name.clone()),
            user_config,
            test_runtime.config(),
        );

        test_runtime
            .set_exporter(exporter)
            .run_test(scenario())
            .run_validation(validation_procedure());
    }

    #[test]
    fn test_exporter_shared() {
        let test_runtime = TestRuntime::new();
        let user_config = Arc::new(NodeUserConfig::new_exporter_config("test_exporter"));
        let exporter = ExporterWrapper::shared(
            TestExporter::new(test_runtime.counters()),
            test_node(test_runtime.config().name.clone()),
            user_config,
            test_runtime.config(),
        );

        test_runtime
            .set_exporter(exporter)
            .run_test(scenario())
            .run_validation(validation_procedure());
    }

    fn make_chan_with_capacity(
        capacity: usize,
    ) -> (
        mpsc::Sender<NodeControlMsg<String>>,
        mpsc::Sender<String>,
        ExporterInbox<String>,
    ) {
        let (control_tx, control_rx) = mpsc::Channel::<NodeControlMsg<String>>::new(capacity);
        let (pdata_tx, pdata_rx) = mpsc::Channel::<String>::new(capacity);
        (
            control_tx,
            pdata_tx,
            ExporterInbox::new(
                Receiver::Local(LocalReceiver::mpsc(control_rx)),
                Receiver::Local(LocalReceiver::mpsc(pdata_rx)),
                0,
                Interests::empty(),
            ),
        )
    }

    fn make_chan() -> (
        mpsc::Sender<NodeControlMsg<String>>,
        mpsc::Sender<String>,
        ExporterInbox<String>,
    ) {
        make_chan_with_capacity(10)
    }

    fn make_processor_chan_with_capacity(
        capacity: usize,
    ) -> (
        mpsc::Sender<NodeControlMsg<String>>,
        mpsc::Sender<String>,
        ProcessorInbox<String>,
    ) {
        let (control_tx, control_rx) = mpsc::Channel::<NodeControlMsg<String>>::new(capacity);
        let (pdata_tx, pdata_rx) = mpsc::Channel::<String>::new(capacity);
        (
            control_tx,
            pdata_tx,
            ProcessorInbox::new(
                Receiver::Local(LocalReceiver::mpsc(control_rx)),
                Receiver::Local(LocalReceiver::mpsc(pdata_rx)),
                0,
                Interests::empty(),
            ),
        )
    }

    fn make_processor_chan() -> (
        mpsc::Sender<NodeControlMsg<String>>,
        mpsc::Sender<String>,
        ProcessorInbox<String>,
    ) {
        make_processor_chan_with_capacity(10)
    }

    #[tokio::test]
    async fn test_control_priority() {
        let (control_tx, pdata_tx, mut channel) = make_chan();
        let pdata1 = "pdata1".to_owned();
        let pdata2 = "pdata2".to_owned();

        pdata_tx.send_async(pdata2.clone()).await.unwrap();
        control_tx
            .send_async(NodeControlMsg::Ack(AckMsg::new(pdata1.clone())))
            .await
            .unwrap();

        // Control message should be received first due to bias
        let msg = channel.recv().await.unwrap();
        assert!(matches!(
            msg,
            Message::Control(NodeControlMsg::Ack(ref a)) if *a.accepted == pdata1
        ));

        // Then pdata message
        let msg = channel.recv().await.unwrap();
        assert!(matches!(msg, Message::PData(ref s) if *s == pdata2));
    }

    #[tokio::test]
    async fn test_shutdown_drain() {
        let (control_tx, pdata_tx, mut channel) = make_chan();

        // Pre-load pdata
        pdata_tx.send_async("pdata1".to_string()).await.unwrap();
        pdata_tx.send_async("pdata2".to_string()).await.unwrap();

        // Send shutdown with a deadline
        control_tx
            .send_async(NodeControlMsg::Shutdown {
                deadline: Instant::now().add(Duration::from_millis(100)), // 100ms deadline
                reason: "Test Shutdown".to_string(),
            })
            .await
            .unwrap();

        // Send more pdata after shutdown is sent, but before the receiver likely gets it
        pdata_tx.send_async("pdata3".to_string()).await.unwrap();
        pdata_tx
            .send_async("pdata4_during_drain".to_string())
            .await
            .unwrap();

        // --- Start Receiving ---

        // 1. Should receive pdata1 (drain)
        let msg1 = channel.recv().await.unwrap();
        assert!(matches!(msg1, Message::PData(ref s) if s == "pdata1"));

        // 2. Should receive pdata2 (drain)
        let msg2 = channel.recv().await.unwrap();
        assert!(matches!(msg2, Message::PData(ref s) if s == "pdata2"));

        // 3. Should receive pdata3 (drain)
        let msg3 = channel.recv().await.unwrap();
        assert!(matches!(msg3, Message::PData(ref s) if s == "pdata3"));

        // 4. Should receive pdata4 (drain)
        let msg4 = channel.recv().await.unwrap();
        assert!(matches!(msg4, Message::PData(ref s) if s == "pdata4_during_drain"));

        // Wait for deadline to likely expire
        sleep(Duration::from_millis(120)).await; // Wait longer than deadline

        // Send pdata *after* deadline
        // This might get buffered but shouldn't be received before the shutdown msg
        let _ = pdata_tx
            .send_async("pdata5_after_deadline".to_string())
            .await;

        // 5. Now, should receive the Shutdown message itself
        let msg5 = channel.recv().await.unwrap();
        // println!("msg5 = {:?}", msg5);
        assert!(matches!(
            msg5,
            Message::Control(NodeControlMsg::Shutdown { .. })
        ));

        drop(control_tx);
        drop(pdata_tx); // Close channels

        // 6. Check for RecvError after channels closed
        let msg_err = channel.recv().await;
        assert!(matches!(msg_err, Err(RecvError::Closed)));
    }

    #[tokio::test]
    async fn test_shutdown_drain_pdata_closes() {
        let (control_tx, pdata_tx, mut channel) = make_chan();

        // Pre-load pdata
        pdata_tx.send_async("pdata1".to_string()).await.unwrap();

        // Send shutdown with a long deadline
        control_tx
            .send_async(NodeControlMsg::Shutdown {
                deadline: Instant::now().add(Duration::from_secs(5)), // Long deadline
                reason: "Test Shutdown PData Closes".to_string(),
            })
            .await
            .unwrap();

        sleep(Duration::from_millis(10)).await; // Give the receiver a chance

        // --- Start Receiving ---

        // 1. Should receive pdata1 (drain)
        let msg1 = channel.recv().await.unwrap();
        assert!(matches!(msg1, Message::PData(ref s) if s == "pdata1"));

        // Close the pdata channel during drain
        drop(pdata_tx);

        // 2. Now, should receive the Shutdown message because pdata channel closed
        let msg2 = channel.recv().await.unwrap();
        assert!(matches!(
            msg2,
            Message::Control(NodeControlMsg::Shutdown { .. })
        ));

        drop(control_tx);

        // 3. Check for RecvError after channels closed
        let msg_err = channel.recv().await;
        assert!(matches!(msg_err, Err(RecvError::Closed)));
    }

    #[tokio::test]
    async fn test_immediate_shutdown() {
        let (control_tx, pdata_tx, mut channel) = make_chan();

        pdata_tx.send_async("pdata1".to_string()).await.unwrap();
        control_tx
            .send_async(NodeControlMsg::Shutdown {
                deadline: Instant::now(), // Immediate deadline
                reason: "Immediate Shutdown".to_string(),
            })
            .await
            .unwrap();

        // Should immediately receive the shutdown message, no draining
        let msg1 = channel.recv().await.unwrap();
        assert!(matches!(
            msg1,
            Message::Control(NodeControlMsg::Shutdown { .. })
        ));

        // Pdata should be ignored and the recv method should return Closed
        let msg2 = channel.recv().await;
        assert!(matches!(msg2, Err(RecvError::Closed)));
    }

    /// After Shutdown all later control messages are silently dropped (ignored).
    #[tokio::test]
    async fn test_ignore_ctrl_after_shutdown() {
        let (control_tx, pdata_tx, mut chan) = make_chan();

        control_tx
            .send_async(NodeControlMsg::Shutdown {
                deadline: Instant::now(),
                reason: "ignore_followups".into(),
            })
            .await
            .unwrap();

        let msg = chan.recv().await.unwrap();
        assert!(matches!(
            msg,
            Message::Control(NodeControlMsg::Shutdown { .. })
        ));

        // Send a control message that should fail as the channel has been closed
        // following the shutdown.
        assert!(
            control_tx
                .send_async(NodeControlMsg::Ack(AckMsg::new("99".to_owned())))
                .await
                .is_err()
        );

        // Send a pdata message that should fail as the channel has been closed
        // following the shutdown.
        assert!(pdata_tx.send_async("pdata1".to_owned()).await.is_err());

        // Another recv should report Closed, proving Ack was discarded.
        assert!(matches!(chan.recv().await, Err(RecvError::Closed)));
    }

    /// Immediate shutdown (deadline == 0) returns Shutdown and then behaves Closed.
    #[tokio::test]
    async fn test_immediate_shutdown_closed_afterwards() {
        let (control_tx, _pdata_tx, mut chan) = make_chan();

        control_tx
            .send_async(NodeControlMsg::Shutdown {
                deadline: Instant::now(),
                reason: "now".into(),
            })
            .await
            .unwrap();

        // First recv -> Shutdown
        let first = chan.recv().await.unwrap();
        assert!(matches!(
            first,
            Message::Control(NodeControlMsg::Shutdown { .. })
        ));

        // Second recv -> channel considered closed
        assert!(matches!(chan.recv().await, Err(RecvError::Closed)));
    }

    // ==================== recv_when tests ====================

    /// recv_when(false) blocks pdata, only returns control messages.
    #[tokio::test]
    async fn test_recv_when_false_blocks_pdata() {
        let (control_tx, pdata_tx, mut channel) = make_processor_chan();

        pdata_tx.send_async("pdata1".to_owned()).await.unwrap();
        control_tx
            .send_async(NodeControlMsg::TimerTick {})
            .await
            .unwrap();

        // recv_when(false) should return the control message, not pdata
        let msg = channel.recv_when(false).await.unwrap();
        assert!(matches!(
            msg,
            Message::Control(NodeControlMsg::TimerTick {})
        ));

        // pdata is still in the channel - recv_when(true) should return it
        let msg = channel.recv_when(true).await.unwrap();
        assert!(matches!(msg, Message::PData(ref s) if s == "pdata1"));
    }

    /// recv_when(true) behaves identically to recv().
    #[tokio::test]
    async fn test_recv_when_true_same_as_recv() {
        let (_control_tx, pdata_tx, mut channel) = make_chan();

        pdata_tx.send_async("pdata1".to_owned()).await.unwrap();

        let msg = channel.recv_when(true).await.unwrap();
        assert!(matches!(msg, Message::PData(ref s) if s == "pdata1"));
    }

    /// During shutdown draining, recv_when(true) drains pdata.
    /// When accept_pdata=false, only control messages are delivered during
    /// draining so processors can resolve backpressure via ack/nack.
    #[tokio::test]
    async fn test_recv_when_true_drains_during_shutdown() {
        let (_control_tx, pdata_tx, mut channel) = make_chan();

        // Pre-load pdata
        pdata_tx.send_async("pdata1".to_owned()).await.unwrap();
        pdata_tx.send_async("pdata2".to_owned()).await.unwrap();

        // Send shutdown with deadline
        _control_tx
            .send_async(NodeControlMsg::Shutdown {
                deadline: Instant::now().add(Duration::from_millis(200)),
                reason: "test".to_owned(),
            })
            .await
            .unwrap();

        // With accept_pdata=true, pdata should be drained during shutdown
        let msg1 = channel.recv_when(true).await.unwrap();
        assert!(matches!(msg1, Message::PData(ref s) if s == "pdata1"));

        let msg2 = channel.recv_when(true).await.unwrap();
        assert!(matches!(msg2, Message::PData(ref s) if s == "pdata2"));

        // Close pdata channel to end draining
        drop(pdata_tx);

        // Should get shutdown message
        let msg3 = channel.recv_when(true).await.unwrap();
        assert!(matches!(
            msg3,
            Message::Control(NodeControlMsg::Shutdown { .. })
        ));
    }

    // Preload more control messages than the fairness burst limit plus one pdata.
    // Once admission is open, the channel must force one pdata through instead of
    // letting sustained control traffic starve it indefinitely.
    #[tokio::test]
    async fn test_recv_forces_pdata_after_control_burst() {
        let (control_tx, pdata_tx, mut channel) = make_chan_with_capacity(64);

        pdata_tx.send_async("pdata1".to_owned()).await.unwrap();

        for _ in 0..40 {
            control_tx
                .send_async(NodeControlMsg::TimerTick {})
                .await
                .unwrap();
        }

        for _ in 0..32 {
            let msg = channel.recv().await.unwrap();
            assert!(matches!(
                msg,
                Message::Control(NodeControlMsg::TimerTick {})
            ));
        }

        let msg = channel.recv_when(true).await.unwrap();
        assert!(matches!(msg, Message::PData(ref s) if s == "pdata1"));

        let msg = channel.recv().await.unwrap();
        assert!(matches!(
            msg,
            Message::Control(NodeControlMsg::TimerTick {})
        ));
    }

    // The same bounded-fair rule must hold while draining after Shutdown is latched.
    // Buffered pdata should still be surfaced once the control burst limit is hit
    // when the caller is accepting pdata during shutdown.
    #[tokio::test]
    async fn test_recv_when_true_forces_pdata_after_control_burst_during_shutdown() {
        let (control_tx, pdata_tx, mut channel) = make_chan_with_capacity(64);

        pdata_tx.send_async("pdata1".to_owned()).await.unwrap();

        control_tx
            .send_async(NodeControlMsg::Shutdown {
                deadline: Instant::now().add(Duration::from_millis(200)),
                reason: "test".to_owned(),
            })
            .await
            .unwrap();

        for _ in 0..40 {
            control_tx
                .send_async(NodeControlMsg::TimerTick {})
                .await
                .unwrap();
        }

        for _ in 0..32 {
            let msg = channel.recv_when(true).await.unwrap();
            assert!(matches!(
                msg,
                Message::Control(NodeControlMsg::TimerTick {})
            ));
        }

        let msg = channel.recv_when(true).await.unwrap();
        assert!(matches!(msg, Message::PData(ref s) if s == "pdata1"));
    }

    // Fairness must not punch through closed processor admission.
    // Even after a large control burst, recv_when(false) keeps pdata buffered
    // until the processor explicitly reopens admission.
    #[tokio::test]
    async fn test_recv_when_false_does_not_bypass_admission_after_control_burst() {
        let (control_tx, pdata_tx, mut channel) = make_processor_chan_with_capacity(64);

        pdata_tx.send_async("pdata1".to_owned()).await.unwrap();

        for _ in 0..40 {
            control_tx
                .send_async(NodeControlMsg::TimerTick {})
                .await
                .unwrap();
        }

        for _ in 0..40 {
            let msg = channel.recv_when(false).await.unwrap();
            assert!(matches!(
                msg,
                Message::Control(NodeControlMsg::TimerTick {})
            ));
        }

        let result =
            tokio::time::timeout(Duration::from_millis(50), channel.recv_when(false)).await;
        assert!(
            result.is_err(),
            "recv_when(false) should keep pdata buffered even after the control burst limit"
        );

        let msg = channel.recv_when(true).await.unwrap();
        assert!(matches!(msg, Message::PData(ref s) if s == "pdata1"));
    }

    /// During shutdown draining with accept_pdata=false, pdata is NOT drained.
    /// Instead, control messages are delivered so processors can reduce in-flight
    /// state and reopen capacity.
    #[tokio::test]
    async fn test_processor_recv_when_false_delivers_control_during_shutdown() {
        let (control_tx, pdata_tx, mut channel) = make_processor_chan();

        // Pre-load pdata
        pdata_tx.send_async("pdata1".to_owned()).await.unwrap();

        // Send shutdown with deadline
        control_tx
            .send_async(NodeControlMsg::Shutdown {
                deadline: Instant::now().add(Duration::from_millis(200)),
                reason: "test".to_owned(),
            })
            .await
            .unwrap();

        // Send a control message (simulating an ack/timer that resolves backpressure)
        control_tx
            .send_async(NodeControlMsg::TimerTick {})
            .await
            .unwrap();

        // recv_when(false) during draining should deliver the control message, not pdata
        let msg = channel.recv_when(false).await.unwrap();
        assert!(
            matches!(msg, Message::Control(NodeControlMsg::TimerTick {})),
            "should deliver control message during draining when accept_pdata=false"
        );

        // Now call with accept_pdata=true -- should drain the buffered pdata
        let msg = channel.recv_when(true).await.unwrap();
        assert!(matches!(msg, Message::PData(ref s) if s == "pdata1"));

        // Close pdata channel
        drop(pdata_tx);

        // Should get shutdown
        let msg = channel.recv_when(true).await.unwrap();
        assert!(matches!(
            msg,
            Message::Control(NodeControlMsg::Shutdown { .. })
        ));
    }

    // Exporters own their receive loop, so shutdown draining is allowed to
    // override temporary admission closure and flush the bounded channel backlog
    // before the final Shutdown is returned.
    #[tokio::test]
    async fn test_exporter_recv_when_false_drains_buffered_pdata_during_shutdown() {
        let (control_tx, pdata_tx, mut channel) = make_chan();

        pdata_tx.send_async("pdata1".to_owned()).await.unwrap();

        control_tx
            .send_async(NodeControlMsg::Shutdown {
                deadline: Instant::now().add(Duration::from_millis(200)),
                reason: "test".to_owned(),
            })
            .await
            .unwrap();

        control_tx
            .send_async(NodeControlMsg::TimerTick {})
            .await
            .unwrap();

        let msg = channel.recv_when(false).await.unwrap();
        assert!(
            matches!(msg, Message::Control(NodeControlMsg::TimerTick {})),
            "exporter should still be able to receive control while draining"
        );

        let msg = channel.recv_when(false).await.unwrap();
        assert!(
            matches!(msg, Message::PData(ref s) if s == "pdata1"),
            "exporter should drain buffered pdata during shutdown even when admission is closed"
        );

        drop(pdata_tx);

        let msg = channel.recv_when(false).await.unwrap();
        assert!(matches!(
            msg,
            Message::Control(NodeControlMsg::Shutdown { .. })
        ));
    }

    /// recv_when(false) with only pdata available should not return pdata.
    /// When a control message is then sent, it should be returned.
    #[tokio::test]
    async fn test_recv_when_false_waits_for_control() {
        let (control_tx, pdata_tx, mut channel) = make_processor_chan();

        pdata_tx.send_async("pdata1".to_owned()).await.unwrap();

        // recv_when(false) should not return pdata -- use a timeout to prove it blocks
        let result =
            tokio::time::timeout(Duration::from_millis(50), channel.recv_when(false)).await;
        assert!(result.is_err(), "recv_when(false) should not return pdata");

        // Now send a control message
        control_tx
            .send_async(NodeControlMsg::TimerTick {})
            .await
            .unwrap();

        // recv_when(false) should now return the control message
        let msg = channel.recv_when(false).await.unwrap();
        assert!(matches!(
            msg,
            Message::Control(NodeControlMsg::TimerTick {})
        ));

        // pdata still buffered
        let msg = channel.recv_when(true).await.unwrap();
        assert!(matches!(msg, Message::PData(ref s) if s == "pdata1"));
    }

    /// recv_when(false) detects a closed pdata channel and generates a
    /// synthetic Shutdown instead of blocking forever.
    #[tokio::test]
    async fn test_recv_when_false_detects_pdata_closed() {
        let (control_tx, pdata_tx, mut channel) = make_processor_chan();

        // Close the pdata channel
        drop(pdata_tx);

        // recv_when(false) should detect the closed channel and return
        // a synthetic Shutdown, not block forever.
        let msg = tokio::time::timeout(Duration::from_millis(100), channel.recv_when(false))
            .await
            .expect("recv_when(false) should not block when pdata channel is closed")
            .unwrap();

        assert!(matches!(
            msg,
            Message::Control(NodeControlMsg::Shutdown { .. })
        ));

        drop(control_tx);
    }

    /// recv_when(false) with a closed pdata channel that still has
    /// buffered data should not trigger synthetic shutdown until the
    /// data is drained.
    #[tokio::test]
    async fn test_recv_when_false_closed_with_buffered_data() {
        let (_control_tx, pdata_tx, mut channel) = make_processor_chan();

        pdata_tx.send_async("pdata1".to_owned()).await.unwrap();
        // Close the channel -- data is still buffered
        drop(pdata_tx);

        // recv_when(false) should NOT trigger shutdown because there's
        // still buffered data (is_empty() is false).
        // It should time out waiting for control.
        let result =
            tokio::time::timeout(Duration::from_millis(50), channel.recv_when(false)).await;
        assert!(
            result.is_err(),
            "should block \u{2014} pdata has data, no control available"
        );

        // Now drain the data with recv_when(true)
        let msg = channel.recv_when(true).await.unwrap();
        assert!(matches!(msg, Message::PData(ref s) if s == "pdata1"));

        // Now recv_when(false) should detect closed+empty and return shutdown
        let msg = channel.recv_when(false).await.unwrap();
        assert!(matches!(
            msg,
            Message::Control(NodeControlMsg::Shutdown { .. })
        ));
    }

    /// Helper that creates a processor-style channel over shared (tokio mpsc)
    /// control and pdata receivers.
    fn make_shared_chan() -> (
        tokio::sync::mpsc::Sender<NodeControlMsg<String>>,
        tokio::sync::mpsc::Sender<String>,
        ProcessorInbox<String>,
    ) {
        let (control_tx, control_rx) = tokio::sync::mpsc::channel::<NodeControlMsg<String>>(10);
        let (pdata_tx, pdata_rx) = tokio::sync::mpsc::channel::<String>(10);
        (
            control_tx,
            pdata_tx,
            ProcessorInbox::new(
                Receiver::Shared(SharedReceiver::mpsc(control_rx)),
                Receiver::Shared(SharedReceiver::mpsc(pdata_rx)),
                0,
                Interests::empty(),
            ),
        )
    }

    /// recv_when(false) on a shared channel with an empty but alive pdata channel
    /// should NOT trigger a synthetic shutdown. This is the regression test for the
    /// bug where SharedReceiver::try_recv mapped Empty to Closed.
    #[tokio::test]
    async fn test_recv_when_false_shared_empty_alive_no_shutdown() {
        let (_control_tx, _pdata_tx, mut channel) = make_shared_chan();

        // Channel is empty but sender is alive.
        // recv_when(false) should block waiting for control -- NOT return a shutdown.
        let result =
            tokio::time::timeout(Duration::from_millis(50), channel.recv_when(false)).await;
        assert!(
            result.is_err(),
            "recv_when(false) on empty alive shared channel should block, not trigger shutdown"
        );
    }

    /// recv_when(false) on a shared channel with a closed pdata channel
    /// should correctly detect closure and return a synthetic shutdown.
    #[tokio::test]
    async fn test_recv_when_false_shared_closed_detects_shutdown() {
        let (_control_tx, pdata_tx, mut channel) = make_shared_chan();

        // Close the pdata channel
        drop(pdata_tx);

        // recv_when(false) should detect the closed channel and return a synthetic Shutdown
        let msg = tokio::time::timeout(Duration::from_millis(100), channel.recv_when(false))
            .await
            .expect("recv_when(false) should not block when shared pdata channel is closed")
            .unwrap();

        assert!(matches!(
            msg,
            Message::Control(NodeControlMsg::Shutdown { .. })
        ));
    }

    // -- with_propagation_policy tests ----------------------------------------

    use crate::context_declaration::CompiledHeaderPropagationPolicy as HeaderPropagationPolicy;

    #[test]
    fn test_with_propagation_policy_none_by_default() {
        let test_runtime = TestRuntime::<TestMsg>::new();
        let wrapper = ExporterWrapper::local(
            TestExporter::new(test_runtime.counters()),
            test_node(test_runtime.config().name.clone()),
            Arc::new(NodeUserConfig::new_exporter_config("test")),
            test_runtime.config(),
        );

        match wrapper {
            ExporterWrapper::Local {
                propagation_policy, ..
            } => assert!(propagation_policy.is_none(), "should be None by default"),
            _ => panic!("expected Local variant"),
        }
    }

    #[test]
    fn test_with_propagation_policy_local() {
        let test_runtime = TestRuntime::<TestMsg>::new();
        let wrapper = ExporterWrapper::local(
            TestExporter::new(test_runtime.counters()),
            test_node(test_runtime.config().name.clone()),
            Arc::new(NodeUserConfig::new_exporter_config("test")),
            test_runtime.config(),
        )
        .with_propagation_policy(Some(HeaderPropagationPolicy::default()));

        match wrapper {
            ExporterWrapper::Local {
                propagation_policy, ..
            } => assert!(
                propagation_policy.is_some(),
                "should be set after with_propagation_policy",
            ),
            _ => panic!("expected Local variant"),
        }
    }

    #[test]
    fn test_with_propagation_policy_shared() {
        let test_runtime = TestRuntime::<TestMsg>::new();
        let wrapper = ExporterWrapper::shared(
            TestExporter::new(test_runtime.counters()),
            test_node(test_runtime.config().name.clone()),
            Arc::new(NodeUserConfig::new_exporter_config("test")),
            test_runtime.config(),
        )
        .with_propagation_policy(Some(HeaderPropagationPolicy::default()));

        match wrapper {
            ExporterWrapper::Shared {
                propagation_policy, ..
            } => assert!(
                propagation_policy.is_some(),
                "should be set after with_propagation_policy",
            ),
            _ => panic!("expected Shared variant"),
        }
    }
}