newton-core 0.7.5

//! Shared telemetry initialization for Newton services.

use super::log::{logger_layer_with_writer, LoggerConfig};
use opentelemetry::{global, trace::TracerProvider as _, KeyValue};
use opentelemetry_sdk::{trace::SdkTracerProvider, Resource};
use std::sync::OnceLock;
use tracing_appender::non_blocking::{ErrorCounter, NonBlocking, NonBlockingBuilder, WorkerGuard};
use tracing_subscriber::{
    filter::filter_fn, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer, Registry,
};

const DEFAULT_TRACE_FILTER: &str = "off";

/// Application-specific values used to initialize telemetry.
#[derive(Debug, Clone, Copy)]
pub struct TelemetryConfig {
    service_name: &'static str,
    service_version: &'static str,
    required_trace_directive: Option<&'static str>,
}

impl TelemetryConfig {
    /// Creates telemetry configuration for a service.
    pub const fn new(service_name: &'static str, service_version: &'static str) -> Self {
        Self {
            service_name,
            service_version,
            required_trace_directive: None,
        }
    }

    /// Requires an exact filter directive before OTLP export is enabled.
    ///
    /// Use this when a service exports only a deliberately enabled trace target.
    pub const fn requiring_trace_directive(mut self, directive: &'static str) -> Self {
        self.required_trace_directive = Some(directive);
        self
    }

    fn trace_export_enabled(self, filter: &str) -> bool {
        match self.required_trace_directive {
            Some(required) => filter.split(',').map(str::trim).any(|directive| directive == required),
            None => filter.trim() != DEFAULT_TRACE_FILTER,
        }
    }
}

/// Flushes completed OpenTelemetry spans, then buffered log lines, when the process exits normally.
#[derive(Debug, Default)]
pub struct TelemetryGuard {
    provider: Option<SdkTracerProvider>,
    // Dropped after `drop` shuts the provider down, so its shutdown errors are flushed too.
    _log_writer: Option<WorkerGuard>,
}

/// Logs are formatted on the emitting thread, then handed to one writer thread.
///
/// Writing to stderr directly blocks the calling tokio worker whenever journald falls
/// behind, which stalls every task on that worker. Lossy: a full buffer drops lines
/// rather than blocking a worker.
fn log_writer<W: std::io::Write + Send + 'static>(sink: W) -> (NonBlocking, WorkerGuard) {
    NonBlockingBuilder::default()
        .lossy(true)
        .thread_name("newton-log-writer")
        .finish(sink)
}

/// Drop count of the process's stderr log writer, set once by `init_telemetry`.
static DROPPED_LOG_LINES: OnceLock<ErrorCounter> = OnceLock::new();

fn stderr_log_writer() -> (NonBlocking, WorkerGuard) {
    let (writer, guard) = log_writer(std::io::stderr());
    let _ = DROPPED_LOG_LINES.set(writer.error_counter());
    (writer, guard)
}

/// Log lines the non-blocking writer has dropped because its buffer was full, since start.
///
/// `None` when this process did not initialize logging through [`init_telemetry`], so a
/// caller can leave the metric absent rather than report a false zero.
pub fn dropped_log_lines() -> Option<u64> {
    DROPPED_LOG_LINES.get().map(|counter| counter.dropped_lines() as u64)
}

impl Drop for TelemetryGuard {
    fn drop(&mut self) {
        if let Some(provider) = self.provider.take() {
            if let Err(error) = provider.shutdown() {
                eprintln!("failed to flush OpenTelemetry traces during shutdown: {error}");
            }
        }
    }
}

/// Initializes Newton's standard logger and optional OTLP trace export.
///
/// The exporter uses the SDK batch processor, keeping network export outside
/// instrumented application futures. The returned guard must live until process
/// shutdown so buffered spans can be flushed.
pub fn init_telemetry(logger: LoggerConfig, config: TelemetryConfig) -> eyre::Result<TelemetryGuard> {
    init_telemetry_with_layers(logger, config, Vec::new())
}

/// Same as [`init_telemetry`], plus caller-supplied layers.
///
/// `newton-metric` depends on `newton-core`, so core cannot reference the
/// metrics crate directly. Binaries pass its pipeline-stage metrics layer
/// (`newton_metric::stages::stage_metrics_layer()`) through here instead.
///
/// Extra layers are installed on BOTH paths -- with and without OTLP export --
/// so metrics derived from spans do not silently depend on whether tracing
/// export happens to be configured.
pub fn init_telemetry_with_layers(
    logger: LoggerConfig,
    config: TelemetryConfig,
    extra_layers: Vec<Box<dyn Layer<Registry> + Send + Sync>>,
) -> eyre::Result<TelemetryGuard> {
    let endpoint_configured = std::env::var_os("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT").is_some();
    let trace_filter = std::env::var("OTEL_TRACES_FILTER").unwrap_or_else(|_| DEFAULT_TRACE_FILTER.to_string());

    if !endpoint_configured || !config.trace_export_enabled(&trace_filter) {
        let (writer, log_guard) = stderr_log_writer();
        let mut layers: Vec<Box<dyn Layer<Registry> + Send + Sync>> = vec![logger_layer_with_writer(logger, writer)];
        layers.extend(extra_layers);
        tracing_subscriber::registry()
            .with(layers)
            .try_init()
            .map_err(|error| eyre::eyre!("failed to initialize tracing subscriber: {error}"))?;
        return Ok(TelemetryGuard {
            provider: None,
            _log_writer: Some(log_guard),
        });
    }

    let protocol = std::env::var("OTEL_EXPORTER_OTLP_TRACES_PROTOCOL").unwrap_or_else(|_| "http/protobuf".to_string());
    if protocol != "http/protobuf" {
        eyre::bail!(
            "{} OTLP export supports http/protobuf, but configured protocol is {protocol}",
            config.service_name
        );
    }

    let exporter = opentelemetry_otlp::SpanExporter::builder()
        .with_http()
        .build()
        .map_err(|error| eyre::eyre!("failed to build OTLP trace exporter: {error}"))?;
    let service_name = std::env::var("OTEL_SERVICE_NAME").unwrap_or_else(|_| config.service_name.to_string());
    let instance_id = std::env::var("HOSTNAME").unwrap_or_else(|_| format!("pid-{}", std::process::id()));
    let resource = Resource::builder()
        .with_service_name(service_name)
        .with_attributes([
            KeyValue::new("service.version", config.service_version),
            KeyValue::new("service.instance.id", instance_id),
        ])
        .build();

    // The SDK builder reads the standard sampler, span-limit, and batch-export
    // environment variables.
    let provider = SdkTracerProvider::builder()
        .with_resource(resource)
        .with_batch_exporter(exporter)
        .build();
    let tracer = provider.tracer(config.service_name);
    let otel_layer = tracing_opentelemetry::layer()
        .with_tracer(tracer)
        .with_filter(EnvFilter::new(trace_filter))
        // Logs keep their normal formatting and destination; only spans are
        // exported through OTLP.
        .with_filter(filter_fn(|metadata| metadata.is_span()))
        .boxed();
    let (writer, log_guard) = stderr_log_writer();
    let mut layers: Vec<Box<dyn Layer<Registry> + Send + Sync>> =
        vec![logger_layer_with_writer(logger, writer), otel_layer];
    layers.extend(extra_layers);

    tracing_subscriber::registry()
        .with(layers)
        .try_init()
        .map_err(|error| eyre::eyre!("failed to initialize tracing subscriber: {error}"))?;
    global::set_tracer_provider(provider.clone());

    Ok(TelemetryGuard {
        provider: Some(provider),
        _log_writer: Some(log_guard),
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn default_gate_exports_any_explicitly_enabled_filter() {
        let config = TelemetryConfig::new("test-service", "test-version");

        assert!(!config.trace_export_enabled("off"));
        assert!(config.trace_export_enabled("info"));
        assert!(config.trace_export_enabled("off,test_target=debug"));
    }

    #[test]
    fn directive_gate_requires_an_exact_directive() {
        let config = TelemetryConfig::new("test-service", "test-version")
            .requiring_trace_directive("newton::task_evaluation=debug");

        assert!(config.trace_export_enabled("off,newton::task_evaluation=debug"));
        assert!(!config.trace_export_enabled("debug"));
        assert!(!config.trace_export_enabled("off,newton::task_evaluation=info"));
    }

    /// A sink that blocks every write until released, like journald falling behind.
    #[derive(Clone)]
    struct StalledSink {
        released: std::sync::Arc<(std::sync::Mutex<bool>, std::sync::Condvar)>,
        written: std::sync::Arc<std::sync::Mutex<Vec<u8>>>,
    }

    impl std::io::Write for StalledSink {
        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
            let (lock, released) = &*self.released;
            let _open = released
                .wait_while(lock.lock().expect("sink lock"), |open| !*open)
                .expect("sink lock");
            self.written.lock().expect("sink buffer").extend_from_slice(buf);
            Ok(buf.len())
        }

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

    #[test]
    fn stalled_log_sink_does_not_block_the_emitting_thread() {
        let sink = StalledSink {
            released: Default::default(),
            written: Default::default(),
        };
        let (writer, guard) = log_writer(sink.clone());
        let subscriber = tracing_subscriber::registry().with(logger_layer_with_writer(
            LoggerConfig::default().with_env_filter("info".to_string()),
            writer,
        ));

        let start = std::time::Instant::now();
        tracing::subscriber::with_default(subscriber, || {
            for line in 0..1_000 {
                tracing::info!(line, "while the sink is stalled");
            }
        });
        assert!(
            start.elapsed() < std::time::Duration::from_secs(1),
            "logging blocked on the sink for {:?}",
            start.elapsed()
        );

        let (lock, released) = &*sink.released;
        *lock.lock().expect("sink lock") = true;
        released.notify_all();
        drop(guard);
        let written = String::from_utf8(sink.written.lock().expect("sink buffer").clone()).expect("utf8 logs");
        assert!(written.contains("while the sink is stalled"));
    }
}