otel-arrow-dfe-engine 0.61.0

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

//! Common foundation of all effect handlers.

use crate::Interests;
use crate::completion_emission_metrics::CompletionEmissionMetricsHandle;
#[cfg(any(test, feature = "test-utils"))]
use crate::control::WakeupRevision;
use crate::control::{
    AckMsg, NackMsg, PipelineCompletionMsg, PipelineCompletionMsgSender, RuntimeControlMsg,
    RuntimeCtrlMsgSender, WakeupSlot,
};
use crate::error::Error;
use crate::node::NodeId;
use crate::node_local_scheduler::NodeLocalSchedulerHandle;
use crate::runtime_services::PipelineRuntimeServices;
use crate::{WakeupError, WakeupSetOutcome};
use otel_arrow_dfe_channel::error::SendError;
use otel_arrow_dfe_telemetry::error::Error as TelemetryError;
use otel_arrow_dfe_telemetry::metrics::{MetricSet, MetricSetHandler};
use otel_arrow_dfe_telemetry::output_service::{Frame, OutputService, StreamHandle};
use otel_arrow_dfe_telemetry::reporter::MetricsReporter;
use std::net::SocketAddr;
use std::time::{Duration, Instant};
use tokio::net::{TcpListener, UdpSocket};

/// SourceTagging indicates whether the Context will contain empty source frames.
#[derive(Clone, Copy)]
pub enum SourceTagging {
    /// Disabled means no source node-id will be automatically
    /// inserted for nodes that do not otherwise subscribe to
    /// Ack/Nack.
    Disabled,

    /// Enabled means a source node_id will be automatically entered
    /// by creating a new frame in as messagees are sent.
    Enabled,
}

impl SourceTagging {
    /// Indicates that source tagging is required.
    #[must_use]
    pub const fn enabled(self) -> bool {
        matches!(self, Self::Enabled)
    }
}

/// Common implementation of all effect handlers.
///
/// Note: This implementation is `Send`.
#[derive(Clone)]
pub(crate) struct EffectHandlerCore<PData> {
    pub(crate) node_id: NodeId,
    // ToDo refactor the code to avoid using Option here.
    pub(crate) runtime_ctrl_msg_sender: Option<RuntimeCtrlMsgSender<PData>>,
    pub(crate) pipeline_completion_msg_sender: Option<PipelineCompletionMsgSender<PData>>,
    #[allow(dead_code)]
    // Will be used in the future. ToDo report metrics from channel and messages.
    pub(crate) metrics_reporter: MetricsReporter,
    /// Optional node-scoped metrics for completions routed by this effect handler.
    pub(crate) completion_emission_metrics: Option<CompletionEmissionMetricsHandle>,
    /// The outgoing message source tagging mode.
    pub(crate) source_tag: SourceTagging,
    /// Precomputed node interests derived from metric level.
    node_interests: Interests,
    /// Optional processor-local delayed-resume and wakeup scheduler.
    pub(crate) local_scheduler: Option<NodeLocalSchedulerHandle<PData>>,
    /// Services owned by the containing pipeline runtime.
    pub(crate) runtime_services: PipelineRuntimeServices,
}

impl<PData> EffectHandlerCore<PData> {
    /// Creates a new EffectHandlerCore with node_id and a metrics reporter.
    pub(crate) const fn new(
        node_id: NodeId,
        metrics_reporter: MetricsReporter,
        runtime_services: PipelineRuntimeServices,
    ) -> Self {
        Self {
            node_id,
            runtime_ctrl_msg_sender: None,
            pipeline_completion_msg_sender: None,
            metrics_reporter,
            completion_emission_metrics: None,
            source_tag: SourceTagging::Disabled,
            node_interests: Interests::empty(),
            local_scheduler: None,
            runtime_services,
        }
    }

    /// Sets the runtime control message sender for this effect handler.
    pub fn set_runtime_ctrl_msg_sender(
        &mut self,
        runtime_ctrl_msg_sender: RuntimeCtrlMsgSender<PData>,
    ) {
        self.runtime_ctrl_msg_sender = Some(runtime_ctrl_msg_sender);
    }

    /// Sets the pipeline result message sender for this effect handler.
    pub fn set_pipeline_completion_msg_sender(
        &mut self,
        pipeline_completion_msg_sender: PipelineCompletionMsgSender<PData>,
    ) {
        self.pipeline_completion_msg_sender = Some(pipeline_completion_msg_sender);
    }

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

    /// Sets the optional node-scoped completion-emission metrics handle.
    pub fn set_completion_emission_metrics(
        &mut self,
        completion_emission_metrics: Option<CompletionEmissionMetricsHandle>,
    ) {
        self.completion_emission_metrics = completion_emission_metrics;
    }

    /// Sets the processor-local wakeup scheduler for this effect handler.
    pub(crate) fn set_local_scheduler(&mut self, local_scheduler: NodeLocalSchedulerHandle<PData>) {
        self.local_scheduler = Some(local_scheduler);
    }

    /// Returns outgoing messages source tagging mode.
    #[must_use]
    pub const fn source_tagging(&self) -> SourceTagging {
        self.source_tag
    }

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

    /// Returns the precomputed node interests.
    ///
    /// Includes SOURCE_TAGGING when source tagging is enabled.
    #[must_use]
    pub fn node_interests(&self) -> Interests {
        if self.source_tag.enabled() {
            self.node_interests | Interests::SOURCE_TAGGING
        } else {
            self.node_interests
        }
    }

    /// Returns the id of the node associated with this effect handler.
    #[must_use]
    pub(crate) fn node_id(&self) -> NodeId {
        self.node_id.clone()
    }

    /// Print an info message to the engine's diagnostic stream.
    ///
    /// This method provides a standardized way for all nodes in the pipeline
    /// to output informational messages. It never waits for the console: a
    /// full diagnostic queue drops the message.
    pub(crate) async fn info(&self, message: &str) {
        submit_diagnostic(&OutputService::diagnostics(), message);
    }

    /// 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.
    ///
    /// ToDo: return a std::net::TcpListener instead of a tokio::net::tcp::TcpListener to avoid leaking our current dependency on Tokio.
    pub(crate) fn tcp_listener(
        &self,
        addr: SocketAddr,
        receiver_id: NodeId,
    ) -> Result<TcpListener, Error> {
        // Helper closure to convert errors.
        let into_engine_error = |error: std::io::Error| Error::IoError {
            node: receiver_id.clone(),
            error,
        };

        // Create a SO_REUSEADDR + SO_REUSEPORT listener.
        let sock = socket2::Socket::new(
            match addr {
                SocketAddr::V4(_) => socket2::Domain::IPV4,
                SocketAddr::V6(_) => socket2::Domain::IPV6,
            },
            socket2::Type::STREAM,
            None,
        )
        .map_err(into_engine_error)?;

        // Allows multiple sockets to bind to an address/port combination even if a socket in the
        // TIME_WAIT state currently occupies that combination.
        // Goal: Restarting the server quickly without waiting for the OS to release a port.
        sock.set_reuse_address(true).map_err(into_engine_error)?;
        // Explicitly allows multiple sockets to simultaneously bind and listen to the exact same
        // IP and port. Incoming connections or packets are distributed between the sockets
        // (load balancing).
        // Goal: Load balancing incoming connections.
        // TODO: Investigate adding set_reuse_port support for Windows.
        #[cfg(unix)]
        {
            sock.set_reuse_port(true).map_err(into_engine_error)?;
        }
        sock.set_nonblocking(true).map_err(into_engine_error)?;
        sock.bind(&addr.into()).map_err(into_engine_error)?;
        sock.listen(8192).map_err(into_engine_error)?;

        TcpListener::from_std(sock.into()).map_err(into_engine_error)
    }

    /// Creates a non-blocking UDP socket on the given address with socket options defined by the
    /// pipeline engine implementation. It's important for receiver implementer to create UDP
    /// sockets 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.
    ///
    /// ToDo: return a std::net::UdpSocket instead of a tokio::net::UdpSocket to avoid leaking our current dependency on Tokio.
    #[allow(dead_code)]
    pub(crate) fn udp_socket(
        &self,
        addr: SocketAddr,
        receiver_id: NodeId,
    ) -> Result<UdpSocket, Error> {
        // Helper closure to convert errors.
        let into_engine_error = |error: std::io::Error| Error::IoError {
            node: receiver_id.clone(),
            error,
        };

        // Create a SO_REUSEADDR + SO_REUSEPORT UDP socket.
        let sock = socket2::Socket::new(
            match addr {
                SocketAddr::V4(_) => socket2::Domain::IPV4,
                SocketAddr::V6(_) => socket2::Domain::IPV6,
            },
            socket2::Type::DGRAM,
            None,
        )
        .map_err(into_engine_error)?;

        // Goal: Restarting the server quickly without waiting for the OS to release a port.
        sock.set_reuse_address(true).map_err(into_engine_error)?;
        // Explicitly allows multiple sockets to simultaneously bind to the exact same
        // IP and port. Incoming packets are distributed between the sockets
        // (load balancing).
        // Goal: Load balancing incoming packets.
        // TODO: Investigate adding set_reuse_port support for Windows.
        #[cfg(unix)]
        {
            sock.set_reuse_port(true).map_err(into_engine_error)?;
        }
        sock.set_nonblocking(true).map_err(into_engine_error)?;
        sock.bind(&addr.into()).map_err(into_engine_error)?;

        UdpSocket::from_std(sock.into()).map_err(into_engine_error)
    }

    /// Reports the provided metrics to the engine.
    #[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.metrics_reporter.report(metrics)
    }

    /// Reports processor-local wakeup scheduler metrics, if this processor uses
    /// the local wakeup service.
    pub fn report_local_scheduler_metrics(
        &self,
        metrics_reporter: &mut MetricsReporter,
    ) -> Result<(), TelemetryError> {
        if let Some(scheduler) = &self.local_scheduler {
            scheduler.report_metrics(metrics_reporter)?;
        }
        Ok(())
    }

    /// Re-usable function to send a runtime control message. This returns a reference
    /// to the sender to place in a cancelation, for example.
    async fn send_runtime_ctrl_msg(
        &self,
        msg: RuntimeControlMsg<PData>,
    ) -> Result<RuntimeCtrlMsgSender<PData>, SendError<RuntimeControlMsg<PData>>> {
        let runtime_ctrl_msg_sender = self.runtime_ctrl_msg_sender.clone()
            .expect("[Internal Error] Node request sender not set. This is a bug in the pipeline engine implementation.");
        runtime_ctrl_msg_sender.send(msg).await?;
        Ok(runtime_ctrl_msg_sender)
    }

    /// Re-usable function to send a pipeline result message.
    async fn send_pipeline_completion_msg(
        &self,
        msg: PipelineCompletionMsg<PData>,
    ) -> Result<PipelineCompletionMsgSender<PData>, SendError<PipelineCompletionMsg<PData>>> {
        let pipeline_completion_msg_sender = self.pipeline_completion_msg_sender.clone()
            .expect("[Internal Error] Node return sender not set. This is a bug in the pipeline engine implementation.");
        pipeline_completion_msg_sender.send(msg).await?;
        Ok(pipeline_completion_msg_sender)
    }

    /// 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: The timer can only be started once per node.
    pub async fn start_periodic_timer(
        &self,
        duration: Duration,
    ) -> Result<TimerCancelHandle<PData>, Error> {
        let runtime_ctrl_msg_sender = self
            .send_runtime_ctrl_msg(RuntimeControlMsg::StartTimer {
                node_id: self.node_id.index,
                duration,
            })
            .await
            .map_err(|e| Error::RuntimeMsgError {
                error: e.to_string(),
            })?;

        Ok(TimerCancelHandle {
            node_id: self.node_id.index,
            runtime_ctrl_msg_sender,
        })
    }

    /// Starts a cancellable periodic telemetry collection timer that emits CollectTelemetry on the control channel.
    /// Returns a handle that can be used to cancel the telemetry timer.
    pub async fn start_periodic_telemetry(
        &self,
        duration: Duration,
    ) -> Result<TelemetryTimerCancelHandle<PData>, Error> {
        let runtime_ctrl_msg_sender = self
            .send_runtime_ctrl_msg(RuntimeControlMsg::StartTelemetryTimer {
                node_id: self.node_id.index,
                duration,
            })
            .await
            .map_err(|e| Error::RuntimeMsgError {
                error: e.to_string(),
            })?;

        Ok(TelemetryTimerCancelHandle {
            node_id: self.node_id.clone(),
            runtime_ctrl_msg_sender,
        })
    }

    /// Send an AckMsg to the runtime control manager for context unwinding.
    /// This will skip if there are no frames.
    ///
    /// External callers should use [`ConsumerEffectHandlerExtension::notify_ack`]
    /// and [`ConsumerEffectHandlerExtension::notify_nack`] instead.
    /// Those wrappers stamp timing information required for correct duration
    /// metrics. Direct access is gated behind the
    /// `#[doc(hidden)] pub mod _private` module and the
    /// [`_private::AckNackRouting`] trait so that accidental bypass is unlikely.
    pub async fn route_ack(&self, ack: AckMsg<PData>) -> Result<(), Error>
    where
        PData: crate::Unwindable,
    {
        if ack.accepted.has_frames() {
            self.send_pipeline_completion_msg(PipelineCompletionMsg::DeliverAck { ack })
                .await
                .map(|_| {
                    if let Some(metrics) = &self.completion_emission_metrics {
                        let mut metrics = metrics
                            .lock()
                            .unwrap_or_else(|poisoned| poisoned.into_inner());
                        metrics.record_notify_ack_routed();
                    }
                })
                .map_err(|e| Error::RuntimeMsgError {
                    error: e.to_string(),
                })
        } else {
            Ok(())
        }
    }

    /// Send a NackMsg to the runtime control manager for context unwinding.
    /// Same semantics as `route_ack()`.
    pub async fn route_nack(&self, nack: NackMsg<PData>) -> Result<(), Error>
    where
        PData: crate::Unwindable,
    {
        if nack.refused.has_frames() {
            self.send_pipeline_completion_msg(PipelineCompletionMsg::DeliverNack { nack })
                .await
                .map(|_| {
                    if let Some(metrics) = &self.completion_emission_metrics {
                        let mut metrics = metrics
                            .lock()
                            .unwrap_or_else(|poisoned| poisoned.into_inner());
                        metrics.record_notify_nack_routed();
                    }
                })
                .map_err(|e| Error::RuntimeMsgError {
                    error: e.to_string(),
                })
        } else {
            Ok(())
        }
    }

    /// Requeue retained pdata onto this node later.
    pub fn requeue_later(&self, when: Instant, data: Box<PData>) -> Result<(), PData> {
        self.local_scheduler
            .as_ref()
            // Safety: processor runtime preparation installs the node-local scheduler
            // before processor code receives an effect handler.
            .expect("node-local scheduler not set for processor effect handler")
            .requeue_later(when, data)
            .map_err(|data| *data)
    }

    /// Set or replace a processor-local wakeup.
    ///
    /// Wakeups are keyed by [`WakeupSlot`]. Scheduling the same slot again
    /// replaces the previous due time for that slot and assigns a new
    /// scheduler revision for that slot.
    ///
    /// The returned [`WakeupSetOutcome`] tells the caller whether the slot was
    /// newly inserted or whether an existing live wakeup was replaced, and
    /// carries the accepted revision that will later be returned in
    /// `NodeControlMsg::Wakeup`.
    ///
    /// # Errors
    ///
    /// Returns [`WakeupError::Unsupported`] when the processor runtime did not
    /// enable processor-local wakeups. Returns [`WakeupError::ShuttingDown`]
    /// once processor shutdown has been latched. Returns
    /// [`WakeupError::Capacity`] if the processor has reached its configured
    /// live wakeup-slot capacity.
    pub fn set_wakeup(
        &self,
        slot: WakeupSlot,
        when: Instant,
    ) -> Result<WakeupSetOutcome, WakeupError> {
        self.local_scheduler
            .as_ref()
            .ok_or(WakeupError::Unsupported)?
            .set_wakeup(slot, when)
    }

    /// Cancel a previously scheduled processor-local wakeup.
    ///
    /// Returns `true` when a live wakeup for `slot` was removed. Returns
    /// `false` when the slot was not scheduled or when shutdown has already
    /// been latched for the processor.
    #[must_use]
    pub fn cancel_wakeup(&self, slot: WakeupSlot) -> bool {
        self.local_scheduler
            .as_ref()
            .map(|scheduler| scheduler.cancel_wakeup(slot))
            .unwrap_or(false)
    }

    /// Pop the next wakeup from the local scheduler, regardless of whether
    /// it is due. Returns `None` when no wakeup is scheduled or when local
    /// wakeups are not enabled.
    ///
    /// This is intended for testing, where the inbox loop is not running and
    /// wakeups need to be manually delivered.
    #[cfg(any(test, feature = "test-utils"))]
    pub fn pop_wakeup(&self) -> Option<(WakeupSlot, Instant, WakeupRevision)> {
        self.local_scheduler
            .as_ref()
            .and_then(|scheduler| scheduler.pop_next())
    }

    /// Notifies the runtime control manager that this receiver has completed
    /// ingress drain.
    pub async fn notify_receiver_drained(&self) -> Result<(), Error> {
        self.send_runtime_ctrl_msg(RuntimeControlMsg::ReceiverDrained {
            node_id: self.node_id().index,
        })
        .await
        .map(|_| ())
        .map_err(|e| Error::RuntimeMsgError {
            error: e.to_string(),
        })
    }
}

/// Handle to cancel a running timer.
pub struct TimerCancelHandle<PData> {
    node_id: usize,
    runtime_ctrl_msg_sender: RuntimeCtrlMsgSender<PData>,
}

impl<PData> TimerCancelHandle<PData> {
    /// Cancels the timer.
    pub async fn cancel(self) -> Result<(), SendError<RuntimeControlMsg<PData>>> {
        self.runtime_ctrl_msg_sender
            .send(RuntimeControlMsg::CancelTimer {
                node_id: self.node_id,
            })
            .await
    }
}

/// Handle to cancel a running telemetry timer.
pub struct TelemetryTimerCancelHandle<PData> {
    node_id: NodeId,
    runtime_ctrl_msg_sender: RuntimeCtrlMsgSender<PData>,
}

impl<PData> TelemetryTimerCancelHandle<PData> {
    /// Cancels the telemetry collection timer.
    pub async fn cancel(self) -> Result<(), SendError<RuntimeControlMsg<PData>>> {
        self.runtime_ctrl_msg_sender
            .send(RuntimeControlMsg::CancelTelemetryTimer {
                node_id: self.node_id.index,
                _temp: std::marker::PhantomData,
            })
            .await
    }
}

/// Queues one diagnostic line, one frame per message so the line is written whole.
///
/// A full queue drops the line and counts it in `diagnostics_dropped`, like the
/// self-tracing writer, so console backpressure never stalls pipeline work.
fn submit_diagnostic(stream: &StreamHandle, message: &str) {
    let _ = stream.try_submit(Frame::line(message));
}

#[cfg(test)]
mod tests {
    use super::*;
    use otel_arrow_dfe_telemetry::output_service::{OutputSink, OutputStream, StreamId};
    use std::io;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicBool, Ordering};

    /// Sink whose writes block until the test releases them.
    struct StalledSink(Arc<AtomicBool>);

    impl OutputSink for StalledSink {
        fn write_frame(&mut self, _frame: &[u8]) -> io::Result<()> {
            while self.0.load(Ordering::Acquire) {
                std::thread::sleep(Duration::from_millis(5));
            }
            Ok(())
        }

        fn flush(&mut self) -> io::Result<()> {
            Ok(())
        }
    }

    /// Scenario: a node keeps emitting info messages while the diagnostic writer is stalled
    /// and its one-frame queue is full.
    /// Guarantees: every message returns at once, and the messages that did not fit are
    /// counted as dropped diagnostics, so a stalled stderr reader cannot hold up the node.
    #[test]
    fn info_drops_instead_of_waiting_on_a_full_diagnostic_queue() {
        const MESSAGES: u64 = 32;

        let stalled = Arc::new(AtomicBool::new(true));
        let stream = OutputStream::start(
            StreamId::Stderr,
            1,
            1024 * 1024,
            true,
            Box::new(StalledSink(Arc::clone(&stalled))),
        )
        .expect("writer thread spawns");
        let handle = stream.handle();

        let started = Instant::now();
        for _ in 0..MESSAGES {
            submit_diagnostic(&handle, "diagnostic");
        }
        let elapsed = started.elapsed();

        let stats = stream.stats();
        // One frame can sit in the stalled write and one in the queue; the rest must drop.
        assert!(stats.frames_submitted <= 2);
        assert_eq!(stats.frames_submitted + stats.diagnostics_dropped, MESSAGES);
        assert!(
            elapsed < Duration::from_secs(1),
            "info must not wait for queue capacity"
        );
        stalled.store(false, Ordering::Release);
        let _ = stream.shutdown(Duration::from_secs(5));
    }
}