otel-arrow-dfe-engine 0.61.0

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

//! Set of traits and structures used to implement receivers.
//!
//! A receiver is an ingress node that feeds a pipeline with data from external sources while
//! performing the necessary conversions to produce messages in a format recognized by the rest of
//! downstream pipeline nodes (e.g. OTLP or OTAP message format).
//!
//! A receiver can operate in various ways, including:
//!
//! 1. Listening on a socket to receive push-based telemetry data,
//! 2. Being notified of changes in a local directory (e.g. log file monitoring),
//! 3. Actively scraping an endpoint to retrieve the latest metrics from a system,
//! 4. Or using any other method to receive or extract telemetry data from external sources.
//!
//! # Lifecycle
//!
//! 1. The receiver is instantiated and configured.
//! 2. The `start` method is called, which begins the receiver's operation.
//! 3. The receiver processes both internal control messages and external data.
//! 4. The receiver shuts down when it receives a `Shutdown` control message or encounters a fatal
//!    error.
//!
//! # Thread Safety
//!
//! This implementation is designed for use in both single-threaded and multi-threaded environments.
//! The `Exporter` trait requires the `Send` bound, enabling the use of thread-safe types.
//!
//! # Scalability
//!
//! To ensure scalability, the pipeline engine will start multiple instances of the same pipeline in
//! parallel on different cores, each with its own receiver instance.

use crate::Interests;
use crate::control::{NodeControlMsg, RuntimeCtrlMsgSender};
use crate::effect_handler::{
    EffectHandlerCore, SourceTagging, TelemetryTimerCancelHandle, TimerCancelHandle,
};
use crate::error::{Error, TypedError};
use crate::node::NodeId;
use crate::output_router::OutputRouter;
use crate::runtime_services::{CodecEffectHandler, PipelineRuntimeServices};
use crate::shared::message::{SharedReceiver, SharedSender};
use crate::terminal_state::TerminalState;
use async_trait::async_trait;
use otel_arrow_dfe_channel::error::RecvError;
use otel_arrow_dfe_config::PortName;
use otel_arrow_dfe_config::authorized_identity_policy::AuthorizedIdentityPolicy;
use otel_arrow_dfe_config::transport_headers_policy::CompiledHeaderCapturePolicy;
use otel_arrow_dfe_pdata_codec::CodecService;
use otel_arrow_dfe_telemetry::error::Error as TelemetryError;
use otel_arrow_dfe_telemetry::metrics::{MetricSet, MetricSetHandler};
use otel_arrow_dfe_telemetry::reporter::MetricsReporter;
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use tokio::net::TcpListener;

/// A trait for ingress receivers (Send definition).
///
/// Receivers are responsible for accepting data from external sources and converting
/// it into messages that can be processed by the pipeline.
#[async_trait]
pub trait Receiver<PData> {
    /// Similar to local::receiver::Receiver::start, but operates in a Send context.
    async fn start(
        self: Box<Self>,
        ctrl_chan: ControlChannel<PData>,
        effect_handler: EffectHandler<PData>,
    ) -> Result<TerminalState, Error>;
}

/// A channel for receiving control messages (in a Send environment).
///
/// This structure wraps a receiver end of a channel that carries [`NodeControlMsg`]
/// values used to control the behavior of a receiver at runtime.
pub struct ControlChannel<PData> {
    rx: SharedReceiver<NodeControlMsg<PData>>,
}

impl<PData> ControlChannel<PData> {
    /// Creates a new `ControlChannelShared` with the given receiver.
    #[must_use]
    pub fn new(rx: SharedReceiver<NodeControlMsg<PData>>) -> Self {
        Self { rx }
    }

    /// Asynchronously receives the next control message.
    ///
    /// Note: produced outcome metrics are now recorded by the pipeline
    /// controller via `Frame` entries collected during context
    /// unwinding, not here.
    ///
    /// # Errors
    ///
    /// Returns a [`RecvError`] if the channel is closed.
    pub async fn recv(&mut self) -> Result<NodeControlMsg<PData>, RecvError> {
        self.rx.recv().await
    }
}

/// A `Send` implementation of the EffectHandlerTrait.
#[derive(Clone)]
pub struct EffectHandler<PData> {
    pub(crate) core: EffectHandlerCore<PData>,
    /// Output-port router.
    pub router: OutputRouter<SharedSender<PData>>,
    /// Immutable capture policy shared by request handlers.
    /// `None` disables capture.
    capture_policy: Option<Arc<CompiledHeaderCapturePolicy>>,
    /// Immutable authorized identity policy shared by request handlers.
    authorized_identity_policy: Option<Arc<AuthorizedIdentityPolicy>>,
}

/// Implementation for the `Send` effect handler.
impl<PData> EffectHandler<PData> {
    /// Creates a sendable receiver effect handler.
    #[must_use]
    pub fn new(
        node_id: NodeId,
        msg_senders: HashMap<PortName, SharedSender<PData>>,
        default_port: Option<PortName>,
        runtime_ctrl_msg_sender: RuntimeCtrlMsgSender<PData>,
        metrics_reporter: MetricsReporter,
        runtime_services: PipelineRuntimeServices,
    ) -> Self {
        let mut core = EffectHandlerCore::new(node_id.clone(), metrics_reporter, runtime_services);
        core.set_runtime_ctrl_msg_sender(runtime_ctrl_msg_sender);
        let router = OutputRouter::new(node_id, msg_senders, default_port);
        EffectHandler {
            core,
            router,
            capture_policy: None,
            authorized_identity_policy: None,
        }
    }

    /// Returns the name of the receiver associated with this handler.
    #[must_use]
    pub fn receiver_id(&self) -> NodeId {
        self.core.node_id()
    }

    /// Sets whether outgoing messages need source node tagging.
    pub fn set_source_tagging(&mut self, value: SourceTagging) {
        self.core.set_source_tagging(value);
    }

    /// Returns whether outgoing messages need source node tagging.
    /// This is true when the destination node has multiple input sources.
    #[must_use]
    pub const fn source_tagging(&self) -> SourceTagging {
        self.core.source_tagging()
    }

    /// Returns the list of connected output ports for this receiver.
    #[must_use]
    pub fn connected_ports(&self) -> Vec<PortName> {
        self.router.connected_ports()
    }

    /// Returns the precomputed node interests.
    #[must_use]
    pub fn node_interests(&self) -> Interests {
        self.core.node_interests()
    }

    /// Returns the capture policy.
    ///
    /// `None` disables capture.
    #[must_use]
    pub fn capture_policy(&self) -> Option<&CompiledHeaderCapturePolicy> {
        self.capture_policy.as_deref()
    }

    /// Sets the capture policy for transport header extraction.
    pub fn set_capture_policy(&mut self, policy: Option<CompiledHeaderCapturePolicy>) {
        self.capture_policy = policy.map(Arc::new);
    }

    /// Returns the authorized identity claim projection policy.
    #[must_use]
    pub fn authorized_identity_policy(&self) -> Option<&AuthorizedIdentityPolicy> {
        self.authorized_identity_policy.as_deref()
    }

    /// Sets the authorized identity claim projection policy.
    pub fn set_authorized_identity_policy(&mut self, policy: Option<AuthorizedIdentityPolicy>) {
        self.authorized_identity_policy = policy.map(Arc::new);
    }

    /// Sends a message to the next node(s) in the pipeline.
    ///
    /// # Errors
    ///
    /// Returns an [`Error::ReceiverError`] if the message could not be routed to a port.
    #[inline]
    pub async fn send_message(&self, data: PData) -> Result<(), TypedError<PData>> {
        self.router.send_default(data).await
    }

    /// Attempts to send a message without awaiting.
    ///
    /// Unlike `send_message`, this method returns immediately if the downstream
    /// channel is full, allowing the caller to handle backpressure without awaiting.
    ///
    /// # Errors
    ///
    /// Returns a [`TypedError::ChannelSendError`] containing [`SendError::Full`] if the
    /// channel is full, or [`SendError::Closed`] if the channel is closed.
    /// Returns a [`TypedError::Error`] if no default port is configured.
    #[inline]
    pub fn try_send_message(&self, data: PData) -> Result<(), TypedError<PData>> {
        self.router.try_send_default(data)
    }

    /// Sends a message to a specific named output port.
    #[inline]
    pub async fn send_message_to<P>(&self, port: P, data: PData) -> Result<(), TypedError<PData>>
    where
        P: Into<PortName>,
    {
        self.router.send_to(port, data).await
    }

    /// Attempts to send a message to a specific named output port without awaiting.
    ///
    /// Unlike `send_message_to`, this method returns immediately if the downstream
    /// channel is full, allowing the caller to handle backpressure without awaiting.
    ///
    /// # Errors
    ///
    /// Returns a [`TypedError::ChannelSendError`] containing [`SendError::Full`] if the
    /// channel is full, or [`SendError::Closed`] if the channel is closed.
    /// Returns a [`TypedError::Error`] if the port does not exist.
    #[inline]
    pub fn try_send_message_to<P>(&self, port: P, data: PData) -> Result<(), TypedError<PData>>
    where
        P: Into<PortName>,
    {
        self.router.try_send_to(port, data)
    }

    /// Creates a non-blocking TCP listener on the given address with socket options defined by the
    /// pipeline engine implementation. It's important for receiver implementer to create TCP
    /// listeners via this method to ensure the scalability and the serviceability of the pipeline.
    ///
    /// # Errors
    ///
    /// Returns an [`Error::IoError`] if any step in the process fails.
    pub fn tcp_listener(&self, addr: SocketAddr) -> Result<TcpListener, Error> {
        self.core.tcp_listener(addr, self.receiver_id())
    }

    /// Print an info message to the engine's diagnostic stream (stderr).
    ///
    /// This method provides a standardized way for receivers to output
    /// informational messages. It never waits for the console: a full
    /// diagnostic queue drops the message.
    pub async fn info(&self, message: &str) {
        self.core.info(message).await;
    }

    /// Starts a cancellable periodic timer that emits TimerTick on the control channel.
    /// Returns a handle that can be used to cancel the timer.
    ///
    /// Current limitation: Only one timer can be started by a processor at a time.
    pub async fn start_periodic_timer(
        &self,
        duration: Duration,
    ) -> Result<TimerCancelHandle<PData>, Error> {
        self.core.start_periodic_timer(duration).await
    }

    /// Starts a cancellable periodic telemetry timer that emits CollectTelemetry.
    pub async fn start_periodic_telemetry(
        &self,
        duration: Duration,
    ) -> Result<TelemetryTimerCancelHandle<PData>, Error> {
        self.core.start_periodic_telemetry(duration).await
    }

    /// Notifies the pipeline runtime that this receiver has completed ingress drain.
    pub async fn notify_receiver_drained(&self) -> Result<(), Error> {
        self.core.notify_receiver_drained().await
    }

    /// Reports metrics collected by the receiver.
    #[allow(dead_code)] // Will be used in the future. ToDo report metrics from channel and messages.
    pub(crate) fn report_metrics<M: MetricSetHandler + 'static>(
        &mut self,
        metrics: &mut MetricSet<M>,
    ) -> Result<(), TelemetryError> {
        self.core.report_metrics(metrics)
    }

    // More methods will be added in the future as needed.
}

impl<PData> CodecEffectHandler for EffectHandler<PData> {
    fn codec_service(&self) -> &CodecService {
        self.core.runtime_services.codecs()
    }
}

#[cfg(test)]
mod tests {
    #![allow(missing_docs)]
    use super::*;
    use crate::control::runtime_ctrl_msg_channel;
    use crate::shared::message::SharedSender;
    use crate::testing::test_node;
    use otel_arrow_dfe_channel::error::SendError;
    use otel_arrow_dfe_telemetry::reporter::MetricsReporter;
    use std::collections::HashMap;

    #[test]
    fn effect_handler_try_send_message_success() {
        let (tx, mut rx) = tokio::sync::mpsc::channel::<u64>(10);
        let mut senders = HashMap::new();
        let _ = senders.insert("out".into(), SharedSender::mpsc(tx));

        let (ctrl_tx, _ctrl_rx) = runtime_ctrl_msg_channel(4);
        let (_metrics_rx, metrics_reporter) = MetricsReporter::create_new_and_receiver(1);
        let eh = EffectHandler::new(
            test_node("recv"),
            senders,
            Some("out".into()),
            ctrl_tx,
            metrics_reporter,
            crate::testing::test_pipeline_runtime_services(),
        );

        // Should succeed when channel has capacity
        assert!(eh.try_send_message(42).is_ok());
        assert_eq!(rx.try_recv().unwrap(), 42);
    }

    #[test]
    fn effect_handler_try_send_message_inbox_full() {
        let (tx, _rx) = tokio::sync::mpsc::channel::<u64>(1);
        let mut senders = HashMap::new();
        let _ = senders.insert("out".into(), SharedSender::mpsc(tx));

        let (ctrl_tx, _ctrl_rx) = runtime_ctrl_msg_channel(4);
        let (_metrics_rx, metrics_reporter) = MetricsReporter::create_new_and_receiver(1);
        let eh = EffectHandler::new(
            test_node("recv"),
            senders,
            Some("out".into()),
            ctrl_tx,
            metrics_reporter,
            crate::testing::test_pipeline_runtime_services(),
        );

        // First send should succeed
        assert!(eh.try_send_message(1).is_ok());
        // Second send should fail with Full
        let result = eh.try_send_message(2);
        assert!(matches!(
            result,
            Err(TypedError::ChannelSendError(SendError::Full(2)))
        ));
    }

    #[test]
    fn effect_handler_try_send_message_no_default_sender() {
        let (a_tx, _a_rx) = tokio::sync::mpsc::channel::<u64>(10);
        let (b_tx, _b_rx) = tokio::sync::mpsc::channel::<u64>(10);

        let mut senders = HashMap::new();
        let _ = senders.insert("a".into(), SharedSender::mpsc(a_tx));
        let _ = senders.insert("b".into(), SharedSender::mpsc(b_tx));

        let (ctrl_tx, _ctrl_rx) = runtime_ctrl_msg_channel(4);
        let (_metrics_rx, metrics_reporter) = MetricsReporter::create_new_and_receiver(1);
        let eh = EffectHandler::new(
            test_node("recv"),
            senders,
            None,
            ctrl_tx,
            metrics_reporter,
            crate::testing::test_pipeline_runtime_services(),
        );

        // Should return configuration error when no default sender
        let result = eh.try_send_message(99);
        assert!(matches!(result, Err(TypedError::Error(_))));
    }

    #[test]
    fn effect_handler_try_send_message_to_success() {
        let (a_tx, mut a_rx) = tokio::sync::mpsc::channel::<u64>(10);
        let (b_tx, mut b_rx) = tokio::sync::mpsc::channel::<u64>(10);

        let mut senders = HashMap::new();
        let _ = senders.insert("a".into(), SharedSender::mpsc(a_tx));
        let _ = senders.insert("b".into(), SharedSender::mpsc(b_tx));

        let (ctrl_tx, _ctrl_rx) = runtime_ctrl_msg_channel(4);
        let (_metrics_rx, metrics_reporter) = MetricsReporter::create_new_and_receiver(1);
        let eh = EffectHandler::new(
            test_node("recv"),
            senders,
            None,
            ctrl_tx,
            metrics_reporter,
            crate::testing::test_pipeline_runtime_services(),
        );

        // Should succeed when sending to a specific port
        assert!(eh.try_send_message_to("b", 42).is_ok());
        assert_eq!(b_rx.try_recv().unwrap(), 42);
        // Port 'a' should not have received anything
        assert!(a_rx.try_recv().is_err());
    }

    #[test]
    fn effect_handler_try_send_message_to_channel_full() {
        let (tx, _rx) = tokio::sync::mpsc::channel::<u64>(1);
        let mut senders = HashMap::new();
        let _ = senders.insert("out".into(), SharedSender::mpsc(tx));

        let (ctrl_tx, _ctrl_rx) = runtime_ctrl_msg_channel(4);
        let (_metrics_rx, metrics_reporter) = MetricsReporter::create_new_and_receiver(1);
        let eh = EffectHandler::new(
            test_node("recv"),
            senders,
            None,
            ctrl_tx,
            metrics_reporter,
            crate::testing::test_pipeline_runtime_services(),
        );

        // First send should succeed
        assert!(eh.try_send_message_to("out", 1).is_ok());
        // Second send should fail with Full
        let result = eh.try_send_message_to("out", 2);
        assert!(matches!(
            result,
            Err(TypedError::ChannelSendError(SendError::Full(2)))
        ));
    }

    #[test]
    fn effect_handler_try_send_message_to_unknown_port() {
        let (tx, _rx) = tokio::sync::mpsc::channel::<u64>(10);
        let mut senders = HashMap::new();
        let _ = senders.insert("out".into(), SharedSender::mpsc(tx));

        let (ctrl_tx, _ctrl_rx) = runtime_ctrl_msg_channel(4);
        let (_metrics_rx, metrics_reporter) = MetricsReporter::create_new_and_receiver(1);
        let eh = EffectHandler::new(
            test_node("recv"),
            senders,
            None,
            ctrl_tx,
            metrics_reporter,
            crate::testing::test_pipeline_runtime_services(),
        );

        // Should return error for unknown port
        let result = eh.try_send_message_to("unknown", 99);
        assert!(matches!(result, Err(TypedError::Error(_))));
    }
}